Fix/concurrent write conn

lava/middleware.go

@@ -1,6 +1,13 @@
 package lava
 
-import "context"
+import (
+	"context"
+	"net/http"
+
+	"google.golang.org/grpc/metadata"
+)
+
+type GrpcGatewayMetadata func(ctx context.Context, req *http.Request, rpcPath string, httpPattern string) metadata.MD
 
 type HandlerFunc func(ctx context.Context, req Request) (Response, error)
 

pkg/wsproxy/websocket_proxy.go

@@ -3,10 +3,12 @@ package wsproxy
 import (
 	"bufio"
 	"bytes"
+	"errors"
 	"io"
 	"net"
 	"net/http"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -170,24 +172,25 @@ func (p *Proxy) proxy(w http.ResponseWriter, r *http.Request) {
 		p.timeWait = timeWait
 	}
 
-	conn, err := upgrade.Upgrade(w, r, responseHeader)
+	conn1, err := upgrade.Upgrade(w, r, responseHeader)
 	if err != nil {
 		log.Warn().Err(err).Msg("error upgrading websocket")
 		return
 	}
-	defer conn.Close()
+	defer conn1.Close()
 
 	ctx, cancelFn := context.WithCancel(context.Background())
 	defer cancelFn()
 
+	conn := WsConn{Conn: conn1, mu: &sync.Mutex{}}
 	conn.SetReadLimit(maxMessageSize)
 	conn.SetPingHandler(func(text string) error {
 		logutil.HandlerErr(conn.SetReadDeadline(time.Now().Add(p.timeWait)))
 
 		log.Info().Str("text", text).Msg("websocket received ping frame")
 		// 不设置 write deadline
 		err := conn.WriteControl(websocket.PongMessage, []byte(text), time.Time{})
-		if err == websocket.ErrCloseSent {
+		if errors.Is(err, websocket.ErrCloseSent) {
 			return nil
 		} else if _, ok := err.(net.Error); ok {
 			return nil
@@ -354,3 +357,32 @@ func (w *inMemoryResponseWriter) CloseNotify() <-chan bool {
 	return w.closed
 }
 func (w *inMemoryResponseWriter) Flush() {}
+
+type WsConn struct {
+	*websocket.Conn
+	mu *sync.Mutex
+}
+
+func (ws WsConn) WritePreparedMessage(pm *websocket.PreparedMessage) error {
+	ws.mu.Lock()
+	defer ws.mu.Unlock()
+	return ws.Conn.WritePreparedMessage(pm)
+}
+
+func (ws WsConn) WriteJSON(v interface{}) error {
+	ws.mu.Lock()
+	defer ws.mu.Unlock()
+	return ws.Conn.WriteJSON(v)
+}
+
+func (ws WsConn) WriteControl(messageType int, data []byte, deadline time.Time) error {
+	ws.mu.Lock()
+	defer ws.mu.Unlock()
+	return ws.Conn.WriteControl(messageType, data, deadline)
+}
+
+func (ws WsConn) WriteMessage(messageType int, data []byte) error {
+	ws.mu.Lock()
+	defer ws.mu.Unlock()
+	return ws.Conn.WriteMessage(messageType, data)
+}

servers/grpcs/server.go

@@ -92,6 +92,7 @@ func (s *serviceImpl) DixInject(
 	log log.Logger,
 	conf *Config,
 	gw []*gateway.Mux,
+	metadataHandlers []lava.GrpcGatewayMetadata,
 ) {
 	s.conf = conf
 	if conf.HttpPort == nil {
@@ -165,6 +166,11 @@ func (s *serviceImpl) DixInject(
 		}))
 	}
 
+	httpServer.Use(func(ctx *fiber.Ctx) error {
+		ctx.SetUserContext(ctx.Context())
+		return ctx.Next()
+	})
+
 	app := fiber.New()
 	app.Group("/debug", httputil.StripPrefix(filepath.Join(conf.BaseUrl, "/debug"), debug.Handler))
 
@@ -239,11 +245,17 @@ func (s *serviceImpl) DixInject(
 			return strings.ToUpper(s), true
 		}),
 		runtime.WithMetadata(func(ctx context.Context, request *http.Request) metadata.MD {
-			path, ok := runtime.HTTPPathPattern(ctx)
-			if !ok {
-				return nil
+			rpcPath, _ := runtime.RPCMethod(ctx)
+			path, _ := runtime.HTTPPathPattern(ctx)
+
+			md := metadata.Pairs("http_path", path, "http_method", request.Method, "http_url", request.URL.Path)
+			for _, h := range metadataHandlers {
+				for k, v := range h(ctx, request, rpcPath, path) {
+					md.Append(k, v...)
+				}
 			}
-			return metadata.Pairs("http_path", path, "http_method", request.Method, "http_url", request.URL.Path)
+
+			return md
 		}),
 		runtime.WithErrorHandler(func(ctx context.Context, mux *runtime.ServeMux, marshal runtime.Marshaler, w http.ResponseWriter, request *http.Request, err error) {
 			md, ok := runtime.ServerMetadataFromContext(ctx)

servers/https/server.go

@@ -19,7 +19,6 @@ import (
 	"github.com/pubgo/funk/stack"
 	"github.com/pubgo/funk/version"
 	"github.com/pubgo/opendoc/opendoc"
-	"github.com/valyala/fasthttp"
 	"google.golang.org/grpc/codes"
 
 	"github.com/pubgo/lava/core/debug"
@@ -51,8 +50,6 @@ func (s *serviceImpl) Run() {
 	defer s.stop()
 	s.start()
 	signal.Wait()
-
-	fasthttp.AcquireArgs()
 }
 
 func (s *serviceImpl) Start() { s.start() }