feat: added context.Context to each handler

feat: small code cleanup

Author: exelban

autobahn/main.go

@@ -12,7 +12,7 @@ func main() {
 	r := http.NewServeMux()
 	r.HandleFunc("/", wsServer.Handler)
 
-	wsServer.On("echo", func(c *websocket.Conn, msg *websocket.Message) {
+	wsServer.On("echo", func(ctx context.Context, c *websocket.Conn, msg *websocket.Message) {
 		_ = c.Emit("echo", msg.Data)
 	})
 

channel.go

@@ -21,14 +21,11 @@ func newChannel(id string) *Channel {
 	}
 
 	go func() {
-		for {
-			select {
-			case conn := <-c.delConn:
-				c.mu.Lock()
-				_ = conn.Close()
-				delete(c.connections, conn)
-				c.mu.Unlock()
-			}
+		for conn := range c.delConn {
+			c.mu.Lock()
+			_ = conn.Close()
+			delete(c.connections, conn)
+			c.mu.Unlock()
 		}
 	}()
 

channel_test.go

@@ -18,7 +18,7 @@ func TestChannel_Add(t *testing.T) {
 
 	ch := wsServer.NewChannel("test-channel-add")
 
-	wsServer.OnConnect(func(c *Conn) {
+	wsServer.OnConnect(func(ctx context.Context, c *Conn) {
 		ch.Add(c)
 		require.Equal(t, 1, ch.Count(), "channel must contain only 1 connection")
 	})
@@ -44,7 +44,7 @@ func TestChannel_Emit(t *testing.T) {
 	messageBytes, err := json.Marshal(_message)
 	require.NoError(t, err)
 
-	wsServer.OnConnect(func(c *Conn) {
+	wsServer.OnConnect(func(ctx context.Context, c *Conn) {
 		ch.Add(c)
 		time.Sleep(300 * time.Millisecond)
 		ch.Emit(_message.Name, _message.Data)
@@ -78,7 +78,7 @@ func TestChannel_Remove(t *testing.T) {
 
 	ch := wsServer.NewChannel("test-channel-add")
 
-	wsServer.OnConnect(func(c *Conn) {
+	wsServer.OnConnect(func(ctx context.Context, c *Conn) {
 		ch.Add(c)
 		require.Equal(t, 1, ch.Count(), "channel must contain only 1 connection")
 		ch.Remove(c)

conn.go

@@ -26,7 +26,8 @@ var pingHeader = ws.Header{
 }
 
 var PingInterval = time.Second * 5
-var TextMessage = false
+
+const TextMessage = false
 
 // ID return an connection identifier (could be not unique)
 func (c *Conn) ID() string {
@@ -43,7 +44,10 @@ func (c *Conn) Emit(name string, data interface{}) error {
 		Data: data,
 	}
 
-	b, _ := json.Marshal(msg)
+	b, err := json.Marshal(msg)
+	if err != nil {
+		return err
+	}
 
 	opCode := ws.OpBinary
 	if TextMessage {
@@ -96,8 +100,7 @@ func (c *Conn) Send(data any) error {
 		Length: int64(len(b)),
 	}
 
-	err := c.Write(h, b)
-	return err
+	return c.Write(h, b)
 }
 
 // Close closing websocket connection.

conn_test.go

@@ -80,7 +80,7 @@ func TestConn_Send_bytes(t *testing.T) {
 	defer shutdown()
 
 	msg := []byte{0, 1, 2, 3, 4, 5, 6, 7}
-	wsServer.OnConnect(func(c *Conn) {
+	wsServer.OnConnect(func(ctx context.Context, c *Conn) {
 		time.Sleep(300 * time.Millisecond)
 		err := c.Send(msg)
 		require.NoError(t, err)
@@ -111,7 +111,7 @@ func TestConn_Send_struct(t *testing.T) {
 	}{
 		Value: "test",
 	}
-	wsServer.OnConnect(func(c *Conn) {
+	wsServer.OnConnect(func(ctx context.Context, c *Conn) {
 		time.Sleep(300 * time.Millisecond)
 		err := c.Send(msg)
 		require.NoError(t, err)

websocket.go

@@ -16,7 +16,7 @@ Echo server:
 
 		r.HandleFunc("/ws", wsServer.Handler)
 
-		wsServer.On("echo", func(c *websocket.Conn, msg *websocket.Message) {
+		wsServer.On("echo", func(ctx context.Context, c *websocket.Conn, msg *websocket.Message) {
 			c.Emit("echo", msg.Data)
 		})
 
@@ -38,7 +38,7 @@ Websocket with group:
 
 		ch := wsServer.NewChannel("test")
 
-		wsServer.OnConnect(func(c *websocket.Conn) {
+		wsServer.OnConnect(func(ctx context.Context, c *websocket.Conn) {
 			ch.Add(c)
 			ch.Emit("connection", []byte("new connection come"))
 		})
@@ -75,9 +75,9 @@ type Server struct {
 
 	delChan []chan *Conn
 
-	onConnect    func(c *Conn)
-	onDisconnect func(c *Conn)
-	onMessage    func(c *Conn, h ws.Header, b []byte)
+	onConnect    func(ctx context.Context, c *Conn)
+	onDisconnect func(ctx context.Context, c *Conn)
+	onMessage    func(ctx context.Context, c *Conn, h ws.Header, b []byte)
 
 	done bool
 	mu   sync.RWMutex
@@ -94,7 +94,7 @@ type Message struct {
 // HandlerFunc is a type for handle function all function which has callback have this struct
 // as first element returns pointer to connection
 // its give opportunity to close connection or emit message to exactly this connection.
-type HandlerFunc func(c *Conn, msg *Message)
+type HandlerFunc func(ctx context.Context, c *Conn, msg *Message)
 
 // New websocket server handler with the provided options.
 func New() *Server {
@@ -104,7 +104,7 @@ func New() *Server {
 		broadcast:   make(chan Message),
 		callbacks:   make(map[string]HandlerFunc),
 	}
-	srv.onMessage = func(c *Conn, h ws.Header, b []byte) {
+	srv.onMessage = func(ctx context.Context, c *Conn, h ws.Header, b []byte) {
 		_ = c.Write(h, b)
 	}
 	return srv
@@ -168,6 +168,7 @@ func (s *Server) Shutdown() error {
 
 // Handler get upgrade connection to RFC 6455 and starting listener for it.
 func (s *Server) Handler(w http.ResponseWriter, r *http.Request) {
+	ctx := r.Context()
 	var params url.Values = nil
 
 	conn, _, _, err := ws.UpgradeHTTP(r, w)
@@ -194,7 +195,7 @@ func (s *Server) Handler(w http.ResponseWriter, r *http.Request) {
 		done:   make(chan bool, 1),
 	}
 	connection.startPing()
-	s.addConn(connection)
+	s.addConn(ctx, connection)
 
 	textPending := false
 
@@ -206,7 +207,7 @@ func (s *Server) Handler(w http.ResponseWriter, r *http.Request) {
 		header, _ := ws.ReadHeader(conn)
 		if err = ws.CheckHeader(header, state); err != nil {
 			log.Printf("drop ws connection: %v", err)
-			s.dropConn(connection)
+			s.dropConn(ctx, connection)
 			break
 		}
 
@@ -263,12 +264,12 @@ func (s *Server) Handler(w http.ResponseWriter, r *http.Request) {
 			if err != nil {
 				log.Printf("drop ws connection: OpClose (%v)", err)
 			}
-			s.dropConn(connection)
+			s.dropConn(ctx, connection)
 			break
 		}
 
 		header.Masked = false
-		if err = s.processMessage(connection, header, payload); err != nil {
+		if err = s.processMessage(ctx, connection, header, payload); err != nil {
 			log.Print(err)
 		}
 	}
@@ -316,21 +317,21 @@ func (s *Server) Channels() []string {
 }
 
 // OnConnect function which will be called when new connections come.
-func (s *Server) OnConnect(f func(c *Conn)) {
+func (s *Server) OnConnect(f func(ctx context.Context, c *Conn)) {
 	s.mu.Lock()
 	s.onConnect = f
 	s.mu.Unlock()
 }
 
 // OnDisconnect function which will be called when new connections come.
-func (s *Server) OnDisconnect(f func(c *Conn)) {
+func (s *Server) OnDisconnect(f func(ctx context.Context, c *Conn)) {
 	s.mu.Lock()
 	s.onDisconnect = f
 	s.mu.Unlock()
 }
 
 // OnMessage handling byte message. This function works as echo by default
-func (s *Server) OnMessage(f func(c *Conn, h ws.Header, b []byte)) {
+func (s *Server) OnMessage(f func(ctx context.Context, c *Conn, h ws.Header, b []byte)) {
 	s.mu.Lock()
 	s.onMessage = f
 	s.mu.Unlock()
@@ -372,14 +373,14 @@ func (s *Server) IsClosed() bool {
 	return s.done
 }
 
-func (s *Server) processMessage(c *Conn, h ws.Header, b []byte) error {
+func (s *Server) processMessage(ctx context.Context, c *Conn, h ws.Header, b []byte) error {
 	if len(b) == 0 {
-		s.onMessage(c, h, b)
+		s.onMessage(ctx, c, h, b)
 		return nil
 	}
 
 	if h.OpCode != ws.OpBinary && h.OpCode != ws.OpText {
-		s.onMessage(c, h, b)
+		s.onMessage(ctx, c, h, b)
 		return nil
 	}
 
@@ -393,30 +394,30 @@ func (s *Server) processMessage(c *Conn, h ws.Header, b []byte) error {
 		if err != nil {
 			return err
 		}
-		s.callbacks[msg.Name](c, &Message{
+		s.callbacks[msg.Name](ctx, c, &Message{
 			Name: msg.Name,
 			Data: buf,
 		})
 		return nil
 	}
-	s.onMessage(c, h, b)
+	s.onMessage(ctx, c, h, b)
 
 	return nil
 }
 
-func (s *Server) addConn(conn *Conn) {
-	if !reflect.ValueOf(s.onConnect).IsNil() {
-		go s.onConnect(conn)
+func (s *Server) addConn(ctx context.Context, conn *Conn) {
+	if s.onConnect != nil {
+		go s.onConnect(ctx, conn)
 	}
 
 	s.mu.Lock()
 	s.connections[conn] = true
 	s.mu.Unlock()
 }
 
-func (s *Server) dropConn(conn *Conn) {
+func (s *Server) dropConn(ctx context.Context, conn *Conn) {
 	if !reflect.ValueOf(s.onDisconnect).IsNil() {
-		go s.onDisconnect(conn)
+		go s.onDisconnect(ctx, conn)
 	}
 
 	go func() {

websocket_test.go

@@ -115,7 +115,7 @@ func TestServer_OnConnect(t *testing.T) {
 	messageBytes, err := json.Marshal(msg)
 	require.NoError(t, err)
 
-	wsServer.OnConnect(func(c *Conn) {
+	wsServer.OnConnect(func(ctx context.Context, c *Conn) {
 		time.Sleep(300 * time.Millisecond)
 		err := c.Emit(msg.Name, msg.Data)
 		require.NoError(t, err)
@@ -155,7 +155,7 @@ func TestServer_OnConnect2(t *testing.T) {
 		Length: int64(len(msg)),
 	}
 
-	wsServer.OnConnect(func(c *Conn) {
+	wsServer.OnConnect(func(ctx context.Context, c *Conn) {
 		time.Sleep(300 * time.Millisecond)
 		err := c.Write(h, msg)
 		require.NoError(t, err)
@@ -189,7 +189,7 @@ func TestServer_OnDisconnect(t *testing.T) {
 		Data: []byte("Hello World"),
 	}
 
-	wsServer.OnDisconnect(func(c *Conn) {
+	wsServer.OnDisconnect(func(ctx context.Context, c *Conn) {
 		time.Sleep(300 * time.Millisecond)
 		_ = c.Emit(msg.Name, msg.Data)
 		done <- true
@@ -232,7 +232,7 @@ func TestServer_OnMessage(t *testing.T) {
 	msg := []byte("Hello from byte array")
 
 	done := make(chan bool, 1)
-	wsServer.OnMessage(func(c *Conn, h ws.Header, b []byte) {
+	wsServer.OnMessage(func(ctx context.Context, c *Conn, h ws.Header, b []byte) {
 		require.Equal(t, msg, b, "response message must be the same as send")
 		done <- true
 	})
@@ -283,7 +283,7 @@ func TestServer_On(t *testing.T) {
 
 	done := make(chan bool, 1)
 
-	wsServer.On("LoL", func(c *Conn, msg *Message) {
+	wsServer.On("LoL", func(ctx context.Context, c *Conn, msg *Message) {
 		require.Equal(t, _message.Name, msg.Name)
 		var respData dataStruct
 		require.NoError(t, json.Unmarshal(msg.Data, &respData))
@@ -386,7 +386,7 @@ func TestServerListen(t *testing.T) {
 	messageBytes, err := json.Marshal(message)
 	require.NoError(t, err)
 
-	wsServer.On("echo", func(c *Conn, msg *Message) {
+	wsServer.On("echo", func(ctx context.Context, c *Conn, msg *Message) {
 		require.Equal(t, message.Name, msg.Name)
 		var respData string
 		require.NoError(t, json.Unmarshal(msg.Data, &respData))
@@ -495,15 +495,15 @@ func TestServer_ConnectionClose(t *testing.T) {
 	ticker := time.NewTicker(time.Millisecond * 1)
 	done := make(chan bool, 1)
 
-	wsServer.OnConnect(func(c *Conn) {
+	wsServer.OnConnect(func(ctx context.Context, c *Conn) {
 		ch.Add(c)
 		require.Equal(t, 1, ch.Count(), "channel must contain only 1 connection")
 		log.Print("Connected")
 	})
-	wsServer.OnDisconnect(func(c *Conn) {
+	wsServer.OnDisconnect(func(ctx context.Context, c *Conn) {
 		log.Print("Disconnected")
 	})
-	wsServer.On("test", func(c *Conn, msg *Message) {
+	wsServer.On("test", func(ctx context.Context, c *Conn, msg *Message) {
 		log.Printf("message: %s", msg.Name)
 	})