Feat(rpc)/websocket gate

cmd/juno/juno.go

@@ -274,9 +274,10 @@ const (
 	dbCompressionUsage = "Database compression profile. Options: zstd, snappy, minlz. " +
 		"Use zstd for low storage."
 	rpcRequestTimeoutUsage        = "Maximum time for an RPC request to complete."
-	rpcMaxConcurrentRequestsUsage = "Maximum concurrent HTTP RPC requests; 0 disables the limit."
-	rpcMaxRequestQueueUsage       = "Maximum number of HTTP RPC requests to queue after " +
-		"reaching rpc-max-concurrent-requests limit."
+	rpcMaxConcurrentRequestsUsage = "Maximum concurrent RPC requests, over HTTP and websocket " +
+		"together; 0 disables the limit."
+	rpcMaxRequestQueueUsage = "Maximum number of HTTP RPC requests to queue after " +
+		"reaching rpc-max-concurrent-requests limit. Websocket requests are never queued."
 	rpcMaxBatchSizeUsage = "Maximum number of calls in a single batch request. " +
 		"0 disables the limit."
 	rpcMaxBatchResponseSizeUsage = "Size (in MBs) at which a batch stops being processed. " +

jsonrpc/gate.go

@@ -67,7 +67,24 @@ func (g *Gate) Acquire(ctx context.Context) error {
 	}
 }
 
-// Release frees a processing slot previously taken by a successful Acquire.
+// TryAcquire reserves a processing slot without waiting, reporting whether it
+// got one. A failure is counted as a rejection, the same as ErrServerBusy from
+// Acquire.
+func (g *Gate) TryAcquire() bool {
+	g.increaseActiveReq()
+
+	select {
+	case g.sem <- struct{}{}:
+		return true
+	default:
+		g.decreaseActiveReq()
+		g.rejected.Add(1)
+		return false
+	}
+}
+
+// Release frees a processing slot previously taken by a successful Acquire or
+// TryAcquire.
 func (g *Gate) Release() {
 	<-g.sem
 	g.decreaseActiveReq()

jsonrpc/gate_test.go

@@ -134,3 +134,26 @@ func TestGateZeroQueue(t *testing.T) {
 		wg.Wait()
 	})
 }
+
+func TestGateTryAcquire(t *testing.T) {
+	gate := jsonrpc.NewGate(1, 2)
+
+	require.True(t, gate.TryAcquire())
+	assert.Equal(t, 1, gate.Running())
+	assert.Equal(t, 0, gate.Queued())
+	assert.Equal(t, uint64(0), gate.Rejected())
+
+	require.False(t, gate.TryAcquire(), "no free slot, and TryAcquire must not queue")
+	assert.Equal(t, 1, gate.Running())
+	assert.Equal(t, 0, gate.Queued(), "a refused caller leaves nothing behind in the queue")
+	assert.Equal(t, uint64(1), gate.Rejected())
+
+	gate.Release()
+	assert.Equal(t, 0, gate.Running())
+
+	require.True(t, gate.TryAcquire(), "the slot is reusable once released")
+	gate.Release()
+	assert.Equal(t, 0, gate.Running())
+	assert.Equal(t, 0, gate.Queued())
+	assert.Equal(t, uint64(1), gate.Rejected())
+}

jsonrpc/server.go

@@ -31,6 +31,9 @@ const (
 	// ResponseTooLarge marks a call that was not executed because the batch had
 	// already reached its response size limit.
 	ResponseTooLarge = -32003
+	// ServerBusy marks a request that was not executed because the server had no
+	// free capacity
+	ServerBusy = -32004
 )
 
 var (

jsonrpc/websocket.go

@@ -2,6 +2,7 @@
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"io"
 	"net/http"
@@ -20,26 +21,43 @@
 	maxConns            = 2048 // TODO: an arbitrary default number, should be revisited after monitoring
 )
 
+var serverBusyResponse = func() []byte {
+	b, err := json.Marshal(&response{
+		Version: "2.0",
+		Error:   &Error{Code: ServerBusy, Message: ErrServerBusy.Error()},
+	})
+	if err != nil {
+		panic(err)
+	}
+	return b
+}()
+
 type Websocket struct {
-	rpc            *Server
-	logger         log.StructuredLogger
+	rpc    *Server
+	logger log.StructuredLogger
+	// For logging busy warnings without flooding
+	sampledLogger  log.StructuredLogger
 	connParams     *WebsocketConnParams
 	listener       NewRequestListener
 	shutdown       <-chan struct{}
 	requestTimeout time.Duration
+	gate           *Gate
 
 	// Add connection tracking
 	connSem *semaphore.Weighted
 }
 
 func NewWebsocket(rpc *Server, shutdown <-chan struct{}, logger log.StructuredLogger) *Websocket {
+	const busyLogInterval = time.Second
+
 	ws := &Websocket{
-		rpc:        rpc,
-		logger:     logger,
-		connParams: DefaultWebsocketConnParams(),
-		listener:   &SelectiveListener{},
-		shutdown:   shutdown,
-		connSem:    semaphore.NewWeighted(maxConns),
+		rpc:           rpc,
+		logger:        logger,
+		sampledLogger: log.Sampled(logger, busyLogInterval, 1, 0),
+		connParams:    DefaultWebsocketConnParams(),
+		listener:      &SelectiveListener{},
+		shutdown:      shutdown,
+		connSem:       semaphore.NewWeighted(maxConns),
 	}
 
 	return ws
@@ -62,6 +80,12 @@
 	return ws
 }
 
+// WithGate registers a gate
+func (ws *Websocket) WithGate(g *Gate) *Websocket {
+	ws.gate = g
+	return ws
+}
+
 // WithListener registers a NewRequestListener
 func (ws *Websocket) WithListener(listener NewRequestListener) *Websocket {
 	ws.listener = listener
@@ -116,7 +140,7 @@
 			break
 		}
 		ws.listener.OnNewRequest("any")
-		if err = ws.rpc.HandleReadWriter(wsc.ctx, ws.requestTimeout, wsc); err != nil {
+		if err = ws.handleMessage(wsc); err != nil {
 			break
 		}
 		// From websocket docs: "Read to EOF otherwise connection will hang."
@@ -146,6 +170,27 @@
 	}
 }
 
+func (ws *Websocket) logServerBusy() {
+	ws.sampledLogger.Warn("Rejected websocket RPC request: server is busy",
+		zap.Int("running", ws.gate.Running()),
+		zap.Int("queued", ws.gate.Queued()),
+		zap.Uint64("rejected", ws.gate.Rejected()),
+	)
+}
+
+func (ws *Websocket) handleMessage(wsc *websocketConn) error {
+	if ws.gate != nil {
+		if !ws.gate.TryAcquire() {
+			ws.logServerBusy()
+			_, err := wsc.Write(serverBusyResponse)
+			return err
+		}
+		defer ws.gate.Release()
+	}
+
+	return ws.rpc.HandleReadWriter(wsc.ctx, ws.requestTimeout, wsc)
+}
+
 type WebsocketConnParams struct {
 	// Maximum message size allowed.
 	ReadLimit int64

jsonrpc/websocket_test.go

@@ -296,3 +296,59 @@ func TestWebsocketConnectionLimit(t *testing.T) {
 	require.Equal(t, http.StatusSwitchingProtocols, resp4.StatusCode)
 	require.NoError(t, conn4.Close(websocket.StatusNormalClosure, ""))
 }
+
+func TestWebsocketGateRejectsWhenBusy(t *testing.T) {
+	started := make(chan struct{})
+	release := make(chan struct{})
+	block := jsonrpc.Method{
+		Name: "test_block",
+		Handler: func(ctx context.Context) (int, *jsonrpc.Error) {
+			close(started)
+			<-release
+			return 0, nil
+		},
+	}
+	echo := jsonrpc.Method{
+		Name:    "test_echo",
+		Params:  []jsonrpc.Parameter{{Name: "msg"}},
+		Handler: func(msg string) (string, *jsonrpc.Error) { return msg, nil },
+	}
+
+	rpc := jsonrpc.NewServer(1, log.NewNopZapLogger())
+	require.NoError(t, rpc.RegisterMethods(block, echo))
+	gate := jsonrpc.NewGate(1, 10)
+	ws := jsonrpc.NewWebsocket(rpc, nil, log.NewNopZapLogger()).WithGate(gate)
+	srv := httptest.NewServer(ws)
+	t.Cleanup(srv.Close)
+
+	connA, respA, err := websocket.Dial(t.Context(), srv.URL, nil) //nolint:bodyclose // lib closes it
+	require.NoError(t, err)
+	require.Equal(t, http.StatusSwitchingProtocols, respA.StatusCode)
+	defer connA.Close(websocket.StatusNormalClosure, "")
+	require.NoError(t, connA.Write(t.Context(), websocket.MessageText,
+		[]byte(`{"jsonrpc":"2.0","method":"test_block","params":[],"id":1}`)))
+	<-started
+
+	connB, respB, err := websocket.Dial(t.Context(), srv.URL, nil) //nolint:bodyclose // lib closes it
+	require.NoError(t, err)
+	require.Equal(t, http.StatusSwitchingProtocols, respB.StatusCode)
+	defer connB.Close(websocket.StatusNormalClosure, "")
+	require.NoError(t, connB.Write(t.Context(), websocket.MessageText,
+		[]byte(`{"jsonrpc":"2.0","method":"test_echo","params":["hi"],"id":2}`)))
+	_, got, err := connB.Read(t.Context())
+	require.NoError(t, err)
+	assert.Equal(t,
+		`{"jsonrpc":"2.0","error":{"code":-32004,"message":"server busy"},"id":null}`,
+		string(got))
+
+	close(release)
+	_, _, err = connA.Read(t.Context())
+	require.NoError(t, err)
+	require.Eventually(t, func() bool { return gate.Running() == 0 }, time.Second, 5*time.Millisecond)
+
+	require.NoError(t, connB.Write(t.Context(), websocket.MessageText,
+		[]byte(`{"jsonrpc":"2.0","method":"test_echo","params":["hi"],"id":3}`)))
+	_, got, err = connB.Read(t.Context())
+	require.NoError(t, err)
+	assert.Equal(t, `{"jsonrpc":"2.0","result":"hi","id":3}`, string(got))
+}

node/http.go

@@ -103,25 +103,13 @@ func makeRPCOverHTTP(
 	metricsEnabled bool,
 	corsEnabled bool,
 	rpcRequestTimeout time.Duration,
-	maxConcurrentRequests uint,
-	maxRequestQueue uint,
+	gate *jsonrpc.Gate,
 ) *httpService {
 	var listener jsonrpc.NewRequestListener
 	if metricsEnabled {
 		listener = makeHTTPMetrics()
 	}
 
-	// A single gate shared across all RPC servers (v8/v9/v10) so the limit
-	// protects the whole process, not each version independently. Disabled when
-	// maxConcurrentRequests is 0.
-	var gate *jsonrpc.Gate
-	if maxConcurrentRequests > 0 {
-		gate = jsonrpc.NewGate(maxConcurrentRequests, uint64(maxRequestQueue))
-		if metricsEnabled {
-			makeHTTPGateMetrics(gate)
-		}
-	}
-
 	mux := http.NewServeMux()
 	for path, server := range servers {
 		httpHandler := jsonrpc.NewHTTP(server, logger).
@@ -156,6 +144,7 @@ func makeRPCOverWebsocket(
 	metricsEnabled bool,
 	corsEnabled bool,
 	rpcRequestTimeout time.Duration,
+	gate *jsonrpc.Gate,
 ) *httpService {
 	var listener jsonrpc.NewRequestListener
 	if metricsEnabled {
@@ -167,7 +156,8 @@ func makeRPCOverWebsocket(
 	mux := http.NewServeMux()
 	for path, server := range servers {
 		wsHandler := jsonrpc.NewWebsocket(server, shutdown, logger).
-			WithRequestTimeout(rpcRequestTimeout)
+			WithRequestTimeout(rpcRequestTimeout).
+			WithGate(gate)
 		if listener != nil {
 			wsHandler = wsHandler.WithListener(listener)
 		}

node/metrics.go

@@ -26,6 +26,7 @@
 	namespaceSync   = "sync"
 	namespacePruner = "pruner"
 	subsystemHTTP   = "http"
+	subsystemGate   = "gate"
 )
 
 func makeDBMetrics() db.EventListener {
@@ -107,28 +108,28 @@
 	}
 }
 
-func makeHTTPGateMetrics(gate *jsonrpc.Gate) {
+func makeRPCGateMetrics(gate *jsonrpc.Gate) {
 	active := prometheus.NewGaugeFunc(prometheus.GaugeOpts{
 		Namespace: "rpc",
-		Subsystem: subsystemHTTP,
+		Subsystem: subsystemGate,
 		Name:      "active_requests",
-		Help:      "Number of HTTP RPC requests currently being processed",
+		Help:      "Number of RPC requests currently being processed, over HTTP and websocket",
 	}, func() float64 {
 		return float64(gate.Running())
 	})
 	queued := prometheus.NewGaugeFunc(prometheus.GaugeOpts{
 		Namespace: "rpc",
-		Subsystem: subsystemHTTP,
+		Subsystem: subsystemGate,
 		Name:      "queued_requests",
-		Help:      "Number of HTTP RPC requests waiting for a processing slot",
+		Help:      "Number of HTTP RPC requests waiting for a processing slot.",
 	}, func() float64 {
 		return float64(gate.Queued())
 	})
 	rejected := prometheus.NewCounterFunc(prometheus.CounterOpts{
 		Namespace: "rpc",
-		Subsystem: subsystemHTTP,
+		Subsystem: subsystemGate,
 		Name:      "rejected_requests",
-		Help:      "Total number of HTTP RPC requests rejected because the server was busy",
+		Help:      "Total number of RPC requests rejected because the server was busy.",
 	}, func() float64 {
 		return float64(gate.Rejected())
 	})

node/node.go

@@ -590,6 +590,13 @@
 		"/rpc" + pathV09: jsonrpcServerV09,
 		"/rpc" + pathV08: jsonrpcServerV08,
 	}
+	var rpcGate *jsonrpc.Gate
+	if (cfg.HTTP || cfg.Websocket) && cfg.RPCMaxConcurrentRequests > 0 {
+		rpcGate = jsonrpc.NewGate(cfg.RPCMaxConcurrentRequests, uint64(cfg.RPCMaxRequestQueue))
+		if cfg.Metrics {
+			makeRPCGateMetrics(rpcGate)
+		}
+	}
 	if cfg.HTTP {
 		readinessHandlers := NewReadinessHandlers(chain, syncReader, cfg.ReadinessBlockTolerance)
 		httpHandlers := map[string]http.HandlerFunc{
@@ -609,8 +616,7 @@
 				cfg.Metrics,
 				cfg.RPCCorsEnable,
 				cfg.RPCRequestTimeout,
-				cfg.RPCMaxConcurrentRequests,
-				cfg.RPCMaxRequestQueue,
+				rpcGate,
 			),
 		)
 	}
@@ -625,6 +631,7 @@
 				cfg.Metrics,
 				cfg.RPCCorsEnable,
 				cfg.RPCRequestTimeout,
+				rpcGate,
 			),
 		)
 	}