lesismal
commented
5 months ago 我没有你完整的资源、所以没调试，大概是下面这样子，细节需要自己看下：
package main
import (
"context"
"flag"
"fmt"
"log"
"net/http"
"os"
"os/signal"
"time"
"github.com/lesismal/nbio/nbhttp"
"github.com/lesismal/nbio/nbhttp/websocket"
)
var (
addr = flag.String("addr", ":8080", "http service address")
hub = Hub{
broadcast: make(chan []byte),
register: make(chan *websocket.Conn),
unregister: make(chan *websocket.Conn),
clients: make(map[*websocket.Conn]bool),
}
upgrader = newUpgrader()
)
func main() {
flag.Parse()
go hub.run()
mux := &http.ServeMux{}
mux.HandleFunc("/", serveHome)
mux.HandleFunc("/ws", serveWs)
engine := nbhttp.NewEngine(nbhttp.Config{
Network: "tcp",
Addrs: []string{*addr},
// MaxLoad: 1000000,
ReleaseWebsocketPayload: true,
Handler: mux,
// websocket 单帧最大载荷
MaxWebsocketFramePayloadSize: 512,
})
err := engine.Start()
if err != nil {
fmt.Printf("nbio.Start failed: %v\n", err)
return
}
interrupt := make(chan os.Signal, 1)
signal.Notify(interrupt, os.Interrupt)
<-interrupt
ctx, cancel := context.WithTimeout(context.Background(), time.Second*3)
defer cancel()
engine.Shutdown(ctx)
}
func newUpgrader() *websocket.Upgrader {
u := websocket.NewUpgrader()
// 服务端主动发心跳检测浪费资源
// 让客户端发心跳、服务端只要读超时关闭就可以了，这样不白白浪费服务端定时器、更好
u.KeepaliveTime = 60 * time.Second
u.OnMessage(func(c *websocket.Conn, messageType websocket.MessageType, message []byte) {
hub.broadcast <- message
})
u.OnClose(func(c *websocket.Conn, err error) {
hub.unregister <- c
})
return u
}
func serveHome(w http.ResponseWriter, r *http.Request) {
log.Println(r.URL.Path)
if r.Method != http.MethodGet {
w.WriteHeader(http.StatusMethodNotAllowed)
fmt.Fprintf(w, "Method not allowed")
return
}
http.ServeFile(w, r, "../home.html")
}
func serveWs(w http.ResponseWriter, r *http.Request) {
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
log.Println(err)
return
}
fmt.Println("Upgraded:", conn.RemoteAddr().String())
hub.register <- conn
}
// Hub maintains the set of active clients and broadcasts messages to the clients.
type Hub struct {
// Registered clients.
clients map[*websocket.Conn]bool
// Inbound messages from the clients.
broadcast chan []byte
// Register requests from the clients.
register chan *websocket.Conn
// Unregister requests from clients.
unregister chan *websocket.Conn
}
func (h *Hub) run() {
for {
select {
case client := <-h.register:
h.clients[client] = true
case client := <-h.unregister:
if _, ok := h.clients[client]; ok {
delete(h.clients, client)
}
case message := <-h.broadcast:
for conn := range h.clients {
conn.WriteMessage(websocket.TextMessage, message)
}
}
}
}
ericjing83
guonaihong
github-actions[bot]
您好大佬，
请教下，我想让这个库负责读消息的协程发送任务给工作协程（消费者），就像下面的代码，生产者readPump发送任务给hub，但是，我刚接触nbio，不知道用这个库如何实现下面的代码，向大佬求助： package main
import ( "bytes" "flag" "github.com/fasthttp/websocket" "github.com/valyala/fasthttp" "log" "time" )
const ( // Time allowed to write a message to the peer. writeWait = 10 * time.Second
)
var ( newline = []byte{'\n'} space = []byte{' '} )
var upgrader = websocket.FastHTTPUpgrader{ ReadBufferSize: 1024, WriteBufferSize: 1024, } var addr = flag.String("addr", ":8080", "http service address") var hub = Hub{ broadcast: make(chan []byte), register: make(chan Client), unregister: make(chan Client), clients: make(map[*Client]bool), }
// Client is a middleman between the websocket connection and the hub. type Client struct {
}
// readPump pumps messages from the websocket connection to the hub. // // The application runs readPump in a per-connection goroutine. The application // ensures that there is at most one reader on a connection by executing all // reads from this goroutine. func (c *Client) readPump() { defer func() { hub.unregister <- c c.conn.Close() }() c.conn.SetReadLimit(maxMessageSize) c.conn.SetReadDeadline(time.Now().Add(pongWait)) c.conn.SetPongHandler(func(string) error { c.conn.SetReadDeadline(time.Now().Add(pongWait)); return nil }) for { _, message, err := c.conn.ReadMessage() if err != nil { if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) { log.Printf("error: %v", err) } break } message = bytes.TrimSpace(bytes.Replace(message, newline, space, -1)) hub.broadcast <- message } }
// writePump pumps messages from the hub to the websocket connection. // // A goroutine running writePump is started for each connection. The // application ensures that there is at most one writer to a connection by // executing all writes from this goroutine. func (c *Client) writePump() { ticker := time.NewTicker(pingPeriod) defer func() { ticker.Stop() c.conn.Close() }() for { select { case message, ok := <-c.send: c.conn.SetWriteDeadline(time.Now().Add(writeWait)) if !ok { // The hub closed the channel. c.conn.WriteMessage(websocket.CloseMessage, []byte{}) return }
}
// serveWs handles websocket requests from the peer. func serveWs(ctx fasthttp.RequestCtx) { err := upgrader.Upgrade(ctx, func(conn websocket.Conn) { client := &Client{conn: conn, send: make(chan []byte, 256)} hub.register <- client
}
// Hub maintains the set of active clients and broadcasts messages to the clients. type Hub struct { // Registered clients. clients map[*Client]bool
}
func (h *Hub) run() { for { select { case client := <-h.register: h.clients[client] = true case client := <-h.unregister: if _, ok := h.clients[client]; ok { delete(h.clients, client) close(client.send) } case message := <-h.broadcast: for client := range h.clients { select { case client.send <- message: default: close(client.send) delete(h.clients, client) } } } } }
func serveHome(ctx *fasthttp.RequestCtx) { log.Println(string(ctx.Path()))
}
func main() { flag.Parse() go hub.run()
}