diff --git a/jsonrpc/server.go b/jsonrpc/server.go index 4cc1a5672e..4d9b7630e9 100644 --- a/jsonrpc/server.go +++ b/jsonrpc/server.go @@ -547,7 +547,7 @@ func (s *Server) handleRequest(ctx context.Context, req *Request) (*response, ht errorIndex := 1 if len(tuple) == 3 { errorIndex = 2 - header = (tuple[1].Interface()).(http.Header) + header = tuple[1].Interface().(http.Header) } if errAny := tuple[errorIndex].Interface(); !utils.IsNil(errAny) { diff --git a/jsonrpc/websocket.go b/jsonrpc/websocket.go index 35fdd82403..b3cdaecfb6 100644 --- a/jsonrpc/websocket.go +++ b/jsonrpc/websocket.go @@ -2,6 +2,7 @@ package jsonrpc import ( "context" + "encoding/json" "errors" "io" "net/http" @@ -20,6 +21,17 @@ 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 @@ -27,6 +39,7 @@ type Websocket struct { 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 diff --git a/jsonrpc/websocket_test.go b/jsonrpc/websocket_test.go index 9d69e5021f..289f31b215 100644 --- a/jsonrpc/websocket_test.go +++ b/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)) +} diff --git a/node/http.go b/node/http.go index fc4aeeeb6b..0a9cbeca0f 100644 --- a/node/http.go +++ b/node/http.go @@ -102,25 +102,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). @@ -155,6 +143,7 @@ func makeRPCOverWebsocket( metricsEnabled bool, corsEnabled bool, rpcRequestTimeout time.Duration, + gate *jsonrpc.Gate, ) *httpService { var listener jsonrpc.NewRequestListener if metricsEnabled { @@ -166,7 +155,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) } diff --git a/node/node.go b/node/node.go index d3602c127d..a47fc10cdb 100644 --- a/node/node.go +++ b/node/node.go @@ -543,6 +543,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, synchronizer, cfg.ReadinessBlockTolerance) httpHandlers := map[string]http.HandlerFunc{ @@ -561,8 +568,7 @@ func New(cfg *Config, version string, logLevel *log.Level) (*Node, error) { cfg.Metrics, cfg.RPCCorsEnable, cfg.RPCRequestTimeout, - cfg.RPCMaxConcurrentRequests, - cfg.RPCMaxRequestQueue, + rpcGate, ), ) } @@ -577,6 +583,7 @@ func New(cfg *Config, version string, logLevel *log.Level) (*Node, error) { cfg.Metrics, cfg.RPCCorsEnable, cfg.RPCRequestTimeout, + rpcGate, ), ) } diff --git a/rpc/handlers.go b/rpc/handlers.go index c2d69b0d9b..943a56d2c0 100644 --- a/rpc/handlers.go +++ b/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, diff --git a/rpc/rpccore/rpccore.go b/rpc/rpccore/rpccore.go index 6b92f6bdfe..4624f18fe5 100644 --- a/rpc/rpccore/rpccore.go +++ b/rpc/rpccore/rpccore.go @@ -19,6 +19,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 @@ -95,4 +97,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"} ) diff --git a/rpc/v10/handlers.go b/rpc/v10/handlers.go index 13d131981a..bd7445aeed 100644 --- a/rpc/v10/handlers.go +++ b/rpc/v10/handlers.go @@ -22,6 +22,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 { @@ -39,8 +40,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 // todo(rdr): why do we have the `TraceCacheKey` type and why it feels uncomfortable // to use. It makes no sense, why not use `Felt` or `Hash` directly? @@ -93,6 +95,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 diff --git a/rpc/v10/subscriptions.go b/rpc/v10/subscriptions.go index 70c41aa1f7..a05eaf397a 100644 --- a/rpc/v10/subscriptions.go +++ b/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) diff --git a/rpc/v10/subscriptions_test.go b/rpc/v10/subscriptions_test.go index 9dc13c73c7..8185a5d1c9 100644 --- a/rpc/v10/subscriptions_test.go +++ b/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 @@ -1217,6 +1218,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) diff --git a/rpc/v8/handlers.go b/rpc/v8/handlers.go index ee93366567..8db885858f 100644 --- a/rpc/v8/handlers.go +++ b/rpc/v8/handlers.go @@ -22,6 +22,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[rpccore.TraceCacheKey, []TracedBlockTransaction] submittedTransactionsCache *rpccore.TransactionCache @@ -89,6 +91,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 diff --git a/rpc/v8/subscriptions.go b/rpc/v8/subscriptions.go index 0756ec529b..e9eb75c77d 100644 --- a/rpc/v8/subscriptions.go +++ b/rpc/v8/subscriptions.go @@ -118,6 +118,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()) @@ -132,6 +135,9 @@ func (h *Handler) subscribe( l1HeadSub, l1HeadRecv := getSubscription(subscriber.onL1Head, h.l1Heads) sub.wg.Go(func() { + if h.subscriptionLimiter != nil { + defer h.subscriptionLimiter.Release(1) + } defer func() { h.unsubscribe(sub, id) unsubscribeFeedSubscription(reorgSub) diff --git a/rpc/v8/subscriptions_test.go b/rpc/v8/subscriptions_test.go index f9cc17b33e..f70fc1a00b 100644 --- a/rpc/v8/subscriptions_test.go +++ b/rpc/v8/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" ) var emptyCommitments = core.BlockCommitments{} @@ -545,6 +546,80 @@ func TestSubscribeNewHeads(t *testing.T) { }) } +// TestSubscribeNewHeadsRespectsLimit checks the subscription limiter rejects a +// subscribe once the slot budget is exhausted and reclaims the slot on unsubscribe. +// The acquire/release logic lives in the shared subscribe(), so covering it via +// SubscribeNewHeads exercises the same path used by every Subscribe* entrypoint. +func TestSubscribeNewHeadsRespectsLimit(t *testing.T) { + logger := log.NewNopZapLogger() + + mockCtrl := gomock.NewController(t) + t.Cleanup(mockCtrl.Finish) + mockChain := mocks.NewMockReader(mockCtrl) + + handler := New(mockChain, nil, nil, logger). + WithSubscriptionLimiter(semaphore.NewWeighted(1)) + + // Every subscribe resolves the range (HeadsHeader); the two that acquire a + // slot each send the latest header as the single historical head. + mockChain.EXPECT().HeadsHeader().Return(&core.Header{}, nil).Times(3) + + // First subscription takes the only slot. + server1, client1 := net.Pipe() + t.Cleanup(func() { + require.NoError(t, server1.Close()) + require.NoError(t, client1.Close()) + }) + ctx1, cancel1 := context.WithCancel(t.Context()) + t.Cleanup(cancel1) + conn1 := &fakeConn{w: server1, ctx: ctx1} + id1, rpcErr := handler.SubscribeNewHeads( + context.WithValue(t.Context(), jsonrpc.ConnKey{}, conn1), nil, + ) + require.Nil(t, rpcErr) + // Drain the historical header so the goroutine parks; otherwise Unsubscribe's + // wg.Wait() blocks on it still writing to the unbuffered pipe. + _, err := client1.Read(make([]byte, db.Megabyte)) + require.NoError(t, err) + + // Second subscription is rejected while the slot is held. + server2, client2 := net.Pipe() + t.Cleanup(func() { + require.NoError(t, server2.Close()) + require.NoError(t, client2.Close()) + }) + conn2 := &fakeConn{w: server2, ctx: t.Context()} + id2, rpcErr := handler.SubscribeNewHeads( + context.WithValue(t.Context(), jsonrpc.ConnKey{}, conn2), nil, + ) + assert.Zero(t, id2) + assert.Equal(t, rpccore.ErrTooManySubscriptions, rpcErr) + + // Unsubscribing the first frees the slot. + ok, rpcErr := handler.Unsubscribe( + context.WithValue(t.Context(), jsonrpc.ConnKey{}, conn1), string(id1), + ) + require.Nil(t, rpcErr) + require.True(t, ok) + + // Third subscription now acquires the reclaimed slot. + server3, client3 := net.Pipe() + t.Cleanup(func() { + require.NoError(t, server3.Close()) + require.NoError(t, client3.Close()) + }) + ctx3, cancel3 := context.WithCancel(t.Context()) + t.Cleanup(cancel3) + conn3 := &fakeConn{w: server3, ctx: ctx3} + id3, rpcErr := handler.SubscribeNewHeads( + context.WithValue(t.Context(), jsonrpc.ConnKey{}, conn3), nil, + ) + require.Nil(t, rpcErr) + require.NotZero(t, id3) + _, err = client3.Read(make([]byte, db.Megabyte)) + require.NoError(t, err) +} + func TestSubscribeNewHeadsHistorical(t *testing.T) { client := feeder.NewTestClient(t, &networks.Mainnet) gw := adaptfeeder.New(client) diff --git a/rpc/v9/handlers.go b/rpc/v9/handlers.go index 7c22fdccb2..34ccf67d30 100644 --- a/rpc/v9/handlers.go +++ b/rpc/v9/handlers.go @@ -22,6 +22,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 // todo(rdr): why do we have the `TraceCacheKey` type and why it feels uncomfortable // to use. It makes no sense, why not use `Felt` or `Hash` directly? @@ -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 diff --git a/rpc/v9/subscriptions.go b/rpc/v9/subscriptions.go index 37f6268e97..7d09f6c2c1 100644 --- a/rpc/v9/subscriptions.go +++ b/rpc/v9/subscriptions.go @@ -106,6 +106,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()) @@ -125,6 +128,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) diff --git a/rpc/v9/subscriptions_test.go b/rpc/v9/subscriptions_test.go index f4205e5ab8..4e2dc44122 100644 --- a/rpc/v9/subscriptions_test.go +++ b/rpc/v9/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" ) var emptyCommitments = core.BlockCommitments{} @@ -1031,6 +1032,81 @@ func TestSubscribeNewHeads(t *testing.T) { }) } +// TestSubscribeNewHeadsRespectsLimit checks the subscription limiter rejects a +// subscribe once the slot budget is exhausted and reclaims the slot on unsubscribe. +// The acquire/release logic lives in the shared subscribe(), so covering it via +// SubscribeNewHeads exercises the same path used by every Subscribe* entrypoint. +func TestSubscribeNewHeadsRespectsLimit(t *testing.T) { + logger := log.NewNopZapLogger() + + mockCtrl := gomock.NewController(t) + t.Cleanup(mockCtrl.Finish) + mockChain := mocks.NewMockReader(mockCtrl) + + handler := New(mockChain, nil, nil, logger). + WithSubscriptionLimiter(semaphore.NewWeighted(1)) + + // Every subscribe resolves the range (Height); only the two that acquire a + // slot go on to send one historical header (BlockHeaderByNumber). + mockChain.EXPECT().Height().Return(uint64(0), nil).Times(3) + mockChain.EXPECT().BlockHeaderByNumber(uint64(0)).Return(&core.Header{}, nil).Times(2) + + // First subscription takes the only slot. + server1, client1 := net.Pipe() + t.Cleanup(func() { + require.NoError(t, server1.Close()) + require.NoError(t, client1.Close()) + }) + ctx1, cancel1 := context.WithCancel(t.Context()) + t.Cleanup(cancel1) + conn1 := &fakeConn{w: server1, ctx: ctx1} + id1, rpcErr := handler.SubscribeNewHeads( + context.WithValue(t.Context(), jsonrpc.ConnKey{}, conn1), nil, + ) + require.Nil(t, rpcErr) + // Drain the historical header so the goroutine parks; otherwise Unsubscribe's + // wg.Wait() blocks on it still writing to the unbuffered pipe. + _, err := client1.Read(make([]byte, db.Megabyte)) + require.NoError(t, err) + + // Second subscription is rejected while the slot is held. + server2, client2 := net.Pipe() + t.Cleanup(func() { + require.NoError(t, server2.Close()) + require.NoError(t, client2.Close()) + }) + conn2 := &fakeConn{w: server2, ctx: t.Context()} + id2, rpcErr := handler.SubscribeNewHeads( + context.WithValue(t.Context(), jsonrpc.ConnKey{}, conn2), nil, + ) + assert.Zero(t, id2) + assert.Equal(t, rpccore.ErrTooManySubscriptions, rpcErr) + + // Unsubscribing the first frees the slot. + ok, rpcErr := handler.Unsubscribe( + context.WithValue(t.Context(), jsonrpc.ConnKey{}, conn1), string(id1), + ) + require.Nil(t, rpcErr) + require.True(t, ok) + + // Third subscription now acquires the reclaimed slot. + server3, client3 := net.Pipe() + t.Cleanup(func() { + require.NoError(t, server3.Close()) + require.NoError(t, client3.Close()) + }) + ctx3, cancel3 := context.WithCancel(t.Context()) + t.Cleanup(cancel3) + conn3 := &fakeConn{w: server3, ctx: ctx3} + id3, rpcErr := handler.SubscribeNewHeads( + context.WithValue(t.Context(), jsonrpc.ConnKey{}, conn3), nil, + ) + require.Nil(t, rpcErr) + require.NotZero(t, id3) + _, err = client3.Read(make([]byte, db.Megabyte)) + require.NoError(t, err) +} + func TestSubscribeNewHeadsHistorical(t *testing.T) { client := feeder.NewTestClient(t, &networks.Mainnet) gw := adaptfeeder.New(client)