diff --git a/websockets/client.go b/websockets/client.go new file mode 100644 index 00000000..dc6ec2b2 --- /dev/null +++ b/websockets/client.go @@ -0,0 +1,6 @@ +package main + +type Client struct { + id string + messages chan []byte +} diff --git a/websockets/main.go b/websockets/main.go index 77667715..2eead092 100644 --- a/websockets/main.go +++ b/websockets/main.go @@ -1,15 +1,17 @@ package main import ( - "context" "log" "net/http" - "time" websocket "nhooyr.io/websocket" ) +const MAX_MESSAGE_QUEUE = 16 + func main() { + server := NewServer() + http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { c, err := websocket.Accept(w, r, &websocket.AcceptOptions{ InsecureSkipVerify: true, @@ -18,20 +20,23 @@ func main() { log.Println(err) return } - defer c.Close(websocket.StatusInternalError, "Internal error") + defer c.CloseNow() + + client := &Client{ + messages: make(chan []byte, MAX_MESSAGE_QUEUE), + } + server.RegisterClient(client) + defer server.DeleteClient(client) for { - ctx, cancel := context.WithTimeout(context.Background(), time.Minute) - defer cancel() - - t, message, err := c.Read(ctx) + t, message, err := c.Read(r.Context()) if err != nil || t != websocket.MessageText { break } log.Printf("Received %v", message) - err = c.Write(ctx, websocket.MessageText, message) + err = c.Write(r.Context(), websocket.MessageText, message) if err != nil { break } diff --git a/websockets/server.go b/websockets/server.go new file mode 100644 index 00000000..1853ccab --- /dev/null +++ b/websockets/server.go @@ -0,0 +1,28 @@ +package main + +import "sync" + +type Server struct { + clients map[*Client]struct{} + mutex sync.RWMutex +} + +func NewServer() *Server { + return &Server{ + clients: make(map[*Client]struct{}), + } +} + +func (server *Server) RegisterClient(client *Client) { + server.mutex.Lock() + defer server.mutex.Unlock() + + server.clients[client] = struct{}{} +} + +func (server *Server) DeleteClient(client *Client) { + server.mutex.Lock() + defer server.mutex.Unlock() + + delete(server.clients, client) +}