perf(rpc): rpc versions use same pool

Author: infrmtcs

jsonrpc/websocket.go

@@ -2,6 +2,7 @@ package jsonrpc
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"io"
 	"net/http"
@@ -20,13 +21,25 @@ const (
 	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: InternalError, Message: ErrServerBusy.Error()},
+	})
+	if err != nil {
+		panic(err)
+	}
+	return b
+}()
+
 type Websocket struct {
 	rpc            *Server
 	logger         log.StructuredLogger
 	connParams     *WebsocketConnParams
 	listener       NewRequestListener
 	shutdown       <-chan struct{}
 	requestTimeout time.Duration
+	gate           *Gate
 
 	// Add connection tracking
 	connSem *semaphore.Weighted
@@ -62,6 +75,11 @@ func (ws *Websocket) WithRequestTimeout(d time.Duration) *Websocket {
 	return ws
 }
 
+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 +134,7 @@ func (ws *Websocket) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			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 +164,27 @@ func (ws *Websocket) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
+func (ws *Websocket) handleMessage(wsc *websocketConn) error {
+	if ws.gate != nil {
+		acquireCtx := wsc.ctx
+		if ws.requestTimeout > 0 {
+			var cancel context.CancelFunc
+			acquireCtx, cancel = context.WithTimeout(acquireCtx, ws.requestTimeout)
+			defer cancel()
+		}
+		if err := ws.gate.Acquire(acquireCtx); err != nil {
+			if errors.Is(err, context.Canceled) {
+				return err
+			}
+			_, writeErr := wsc.Write(serverBusyResponse)
+			return writeErr
+		}
+		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, 0)
+	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":-32603,"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/node.go

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

rpc/handlers.go

@@ -19,6 +19,7 @@ import (
 	"github.com/NethermindEth/juno/utils/log"
 	"github.com/NethermindEth/juno/vm"
 	"golang.org/x/sync/errgroup"
+	"golang.org/x/sync/semaphore"
 )
 
 const (
@@ -44,9 +45,13 @@ type Handler struct {
 func New(bcReader blockchain.Reader, syncReader sync.Reader, virtualMachine vm.VM, version string,
 	logger log.Logger, network *networks.Network,
 ) *Handler {
-	handlerv8 := rpcv8.New(bcReader, syncReader, virtualMachine, logger)
-	handlerv9 := rpcv9.New(bcReader, syncReader, virtualMachine, logger)
-	handlerv10 := rpcv10.New(bcReader, syncReader, virtualMachine, logger)
+	subscriptionLimiter := semaphore.NewWeighted(rpccore.DefaultMaxSubscriptions)
+	handlerv8 := rpcv8.New(bcReader, syncReader, virtualMachine, logger).
+		WithSubscriptionLimiter(subscriptionLimiter)
+	handlerv9 := rpcv9.New(bcReader, syncReader, virtualMachine, logger).
+		WithSubscriptionLimiter(subscriptionLimiter)
+	handlerv10 := rpcv10.New(bcReader, syncReader, virtualMachine, logger).
+		WithSubscriptionLimiter(subscriptionLimiter)
 
 	return &Handler{
 		rpcv8Handler:  handlerv8,

rpc/rpccore/rpccore.go

@@ -18,6 +18,8 @@ const (
 	MaxBlocksBack                 = 1024
 	EntrypointNotFoundFelt string = "0x454e545259504f494e545f4e4f545f464f554e44"
 	ErrEPSNotFound                = "Entry point EntryPointSelector(%s) not found in contract."
+
+	DefaultMaxSubscriptions int64 = 2048
 )
 
 //go:generate mockgen -destination=../mocks/mock_gateway_handler.go -package=mocks github.com/NethermindEth/juno/rpc/rpccore Gateway
@@ -90,4 +92,5 @@ var (
 
 	// These errors can be only be returned by Juno-specific methods.
 	ErrSubscriptionNotFound = &jsonrpc.Error{Code: 100, Message: "Subscription not found"}
+	ErrTooManySubscriptions = &jsonrpc.Error{Code: 101, Message: "Too many subscriptions"}
 )

rpc/v10/handlers.go

@@ -23,6 +23,7 @@ import (
 	"github.com/NethermindEth/juno/utils/lru"
 	"github.com/NethermindEth/juno/vm"
 	"github.com/sourcegraph/conc"
+	"golang.org/x/sync/semaphore"
 )
 
 type Handler struct {
@@ -40,8 +41,9 @@ type Handler struct {
 	l1Heads                 *feed.Feed[*core.L1Head]
 	receivedTransactionFeed *feed.Feed[core.Transaction]
 
-	idgen         func() string
-	subscriptions stdsync.Map // map[string]*subscription
+	idgen               func() string
+	subscriptions       stdsync.Map // map[string]*subscription
+	subscriptionLimiter *semaphore.Weighted
 
 	blockTraceCache *lru.Cache[felt.Felt, TraceBlockTransactionsResponse]
 	// todo(rdr): Can this cache be genericified and can it be applied to the `blockTraceCache`
@@ -92,6 +94,11 @@ func New(
 	}
 }
 
+func (h *Handler) WithSubscriptionLimiter(limiter *semaphore.Weighted) *Handler {
+	h.subscriptionLimiter = limiter
+	return h
+}
+
 func (h *Handler) WithCompiler(compiler compiler.Compiler) *Handler {
 	h.compiler = compiler
 	return h

rpc/v10/subscriptions.go

@@ -54,6 +54,9 @@ func (h *Handler) subscribe(
 	wsConn jsonrpc.Conn,
 	subscriber subscriber,
 ) (SubscriptionID, *jsonrpc.Error) {
+	if h.subscriptionLimiter != nil && !h.subscriptionLimiter.TryAcquire(1) {
+		return "", rpccore.ErrTooManySubscriptions
+	}
 	id := h.idgen()
 	//nolint:gosec // G118: cancel called in unsubscribe()
 	subscriptionCtx, subscriptionCtxCancel := context.WithCancel(wsConn.Context())
@@ -75,6 +78,9 @@ func (h *Handler) subscribe(
 	)
 
 	sub.wg.Go(func() {
+		if h.subscriptionLimiter != nil {
+			defer h.subscriptionLimiter.Release(1)
+		}
 		defer func() {
 			h.unsubscribe(sub, id)
 			unsubscribeFeedSubscription(reorgSub)

rpc/v10/subscriptions_test.go

@@ -32,6 +32,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/mock/gomock"
+	"golang.org/x/sync/semaphore"
 )
 
 // mustNewChain builds a pre_confirmed ChainReader from statically valid test
@@ -1214,6 +1215,49 @@ func TestSubscribeNewHeads(t *testing.T) {
 	assertNextHead(t, conn, subID, &adaptedHeader3)
 }
 
+func TestSubscribeNewHeadsRespectsLimit(t *testing.T) {
+	logger := log.NewNopZapLogger()
+	client := feeder.NewTestClient(t, &networks.Sepolia)
+	block1, commitments1, stateUpdate1 := GetTestBlockWithCommitments(t, client, 56377)
+	adaptedHeader := AdaptBlockHeader(block1.Header, commitments1, stateUpdate1.StateDiff)
+
+	mockCtrl := gomock.NewController(t)
+	t.Cleanup(mockCtrl.Finish)
+	mockChain := mocks.NewMockReader(mockCtrl)
+
+	handler := New(mockChain, nil, nil, logger).
+		WithSubscriptionLimiter(semaphore.NewWeighted(1))
+
+	mockChain.EXPECT().Height().Return(block1.Number, nil).Times(3)
+	mockChain.EXPECT().BlockHeaderByNumber(block1.Number).Return(block1.Header, nil).Times(2)
+	mockChain.EXPECT().BlockCommitmentsByNumber(block1.Number).Return(commitments1, nil).Times(2)
+	mockChain.EXPECT().StateUpdateByNumber(block1.Number).Return(stateUpdate1, nil).Times(2)
+
+	blockIDLatest := BlockIDLatest()
+
+	subID1, conn1 := createTestNewHeadsWebsocket(t, handler, (*SubscriptionBlockID)(&blockIDLatest))
+	assertNextHead(t, conn1, subID1, &adaptedHeader)
+
+	serverConn, clientConn := net.Pipe()
+	t.Cleanup(func() {
+		require.NoError(t, serverConn.Close())
+		require.NoError(t, clientConn.Close())
+	})
+	rejConn := &fakeConn{Conn: clientConn, w: serverConn}
+	rejCtx := context.WithValue(t.Context(), jsonrpc.ConnKey{}, rejConn)
+	id, rpcErr := handler.SubscribeNewHeads(rejCtx, nil)
+	assert.Zero(t, id)
+	assert.Equal(t, rpccore.ErrTooManySubscriptions, rpcErr)
+
+	unsubCtx := context.WithValue(t.Context(), jsonrpc.ConnKey{}, conn1)
+	ok, rpcErr := handler.Unsubscribe(unsubCtx, string(subID1))
+	require.Nil(t, rpcErr)
+	require.True(t, ok)
+
+	subID3, conn3 := createTestNewHeadsWebsocket(t, handler, (*SubscriptionBlockID)(&blockIDLatest))
+	assertNextHead(t, conn3, subID3, &adaptedHeader)
+}
+
 func TestSubscribeNewHeadsHistorical(t *testing.T) {
 	logger := log.NewNopZapLogger()
 	client := feeder.NewTestClient(t, &networks.Sepolia)