feat: timeout

Author: nugaon

pkg/api/layer2.go

@@ -27,6 +27,14 @@ func (s *Service) layer2WsHandler(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	headers := struct {
+		KeepAlive time.Duration `map:"Swarm-Keep-Alive"`
+	}{}
+	if response := s.mapStructure(r.Header, &headers); response != nil {
+		response("invalid header params", logger, w)
+		return
+	}
+
 	if s.beeMode == DevMode {
 		logger.Warning("layer2 endpoint is disabled in dev mode")
 		jsonhttp.BadRequest(w, errUnsupportedDevNodeOperation)
@@ -46,7 +54,12 @@ func (s *Service) layer2WsHandler(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	pingPeriod := 100 * time.Second //TODO pass it in header
+	pingPeriod := headers.KeepAlive * time.Second
+	if pingPeriod == 0 {
+		pingPeriod = 30 * time.Second
+	}
+
+	logger.Info("pingPeriod", pingPeriod)
 	ctx, cancel := context.WithCancel(context.Background())
 	protocolService := s.l2p2p.GetProtocol(ctx, paths.StreamName)
 	err = s.p2p.AddProtocol(protocolService.Protocol())
@@ -58,7 +71,8 @@ func (s *Service) layer2WsHandler(w http.ResponseWriter, r *http.Request) {
 
 	s.wsWg.Add(1)
 	go func() {
-		layer2.ListeningWs(ctx, conn, layer2.WsOptions{PingPeriod: pingPeriod}, logger, protocolService)
+		layer2.ListeningWs(ctx, conn, layer2.WsOptions{PingPeriod: pingPeriod, Cancel: cancel}, logger, protocolService)
+		_ = conn.Close()
 		s.wsWg.Done()
 		cancel()
 	}()

pkg/layer2/ws.go

@@ -19,6 +19,7 @@ import (
 
 type WsOptions struct {
 	PingPeriod time.Duration
+	Cancel     context.CancelFunc
 }
 
 type actionType uint8
@@ -75,12 +76,12 @@ func ListeningWs(ctx context.Context, conn *websocket.Conn, options WsOptions, l
 			err = conn.SetReadDeadline(time.Now().Add(readDeadline))
 			if err != nil {
 				logger.Debug("L2 ws: set write deadline failed", "error", err)
-				return
+				break
 			}
 			messageType, p, err := conn.ReadMessage()
 			if err != nil {
 				logger.Error(err, "L2 ws read message")
-				return
+				break
 			}
 			respMessageType.Store(uint32(messageType))
 			if messageType == 1 {
@@ -94,6 +95,7 @@ func ListeningWs(ctx context.Context, conn *websocket.Conn, options WsOptions, l
 					conn, err := protocolService.GetConnection(ctx, overlay)
 					if err != nil {
 						logger.Error(err, "L2 get connection")
+						break
 					}
 					err = conn.SendMessage(ctx, msg)
 					if err != nil {
@@ -126,6 +128,7 @@ func ListeningWs(ctx context.Context, conn *websocket.Conn, options WsOptions, l
 				}
 			}
 		}
+		options.Cancel()
 	}()
 
 	defer func() {
@@ -147,7 +150,6 @@ func ListeningWs(ctx context.Context, conn *websocket.Conn, options WsOptions, l
 			}
 			return
 		case <-ticker.C:
-			err = conn.SetWriteDeadline(time.Now().Add(writeDeadline))
 			if err != nil {
 				logger.Debug("L2 ws: set write deadline failed", "error", err)
 				return