Skip to content

feat(rpc): gate WSS requests and bound active subscriptions #3883

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
infrmtcs wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(rpc): gate WSS requests and bound active subscriptions #3883

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion jsonrpc/server.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -547,7 +547,7 @@
errorIndex := 1
if len(tuple) == 3 {
errorIndex = 2
header = (tuple[1].Interface()).(http.Header)
header = tuple[1].Interface().(http.Header)

Check warning on line 550 in jsonrpc/server.go

View check run for this annotation

Codecov / codecov/patch

jsonrpc/server.go#L550 

Added line #L550 was not covered by tests
}

if errAny := tuple[errorIndex].Interface(); !utils.IsNil(errAny) {
Expand Down
41 changes: 40 additions & 1 deletion jsonrpc/websocket.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import (
"context"
"encoding/json"
"errors"
"io"
"net/http"
Expand All @@ -20,13 +21,25 @@
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)

Check warning on line 30 in jsonrpc/websocket.go

View check run for this annotation

Codecov / codecov/patch

jsonrpc/websocket.go#L30 

Added line #L30 was not covered by tests
}
return b
}()
Comment on lines +24 to +33

@github-actions github-actions Bot Jul 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

important — the busy response is always "id":null, so pipelining clients can't correlate it

Because the gate is checked before the frame is parsed (HandleReadWriter is what reads wsc.r), this canned response carries no request id. Websocket clients routinely have several requests in flight on one connection and match responses by id; an id:null error is unattributable, so a client typically either drops it or fails the wrong request. The HTTP path doesn't have this problem — it answers with a 503 at the transport level, which is unambiguous.

Related: a client that sent a notification (no id) now receives an unsolicited response frame, which is a JSON-RPC violation.

Options: read the frame into a bounded buffer and decode just the id before gating (cost: one extra copy bounded by ReadLimit), or close the connection with a websocket close code instead of a fake response. If the current behaviour is a deliberate trade-off, please note it in a comment on serverBusyResponse.

Separately, InternalError (-32603) is a poor fit semantically — nothing failed internally, the server is shedding load. A dedicated code (or ErrServerBusy mapped to one) would let clients implement backoff instead of treating it as a bug.


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
Expand Down Expand Up @@ -62,6 +75,11 @@
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
Expand Down Expand Up @@ -116,7 +134,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."
Expand Down Expand Up @@ -146,6 +164,27 @@
}
}

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()

Check warning on line 173 in jsonrpc/websocket.go

View check run for this annotation

Codecov / codecov/patch

jsonrpc/websocket.go#L171-L173 

Added lines #L171 - L173 were not covered by tests
}
if err := ws.gate.Acquire(acquireCtx); err != nil {
if errors.Is(err, context.Canceled) {
return err

Check warning on line 177 in jsonrpc/websocket.go

View check run for this annotation

Codecov / codecov/patch

jsonrpc/websocket.go#L177 

Added line #L177 was not covered by tests
}
_, writeErr := wsc.Write(serverBusyResponse)
return writeErr
}
Comment on lines +175 to +181

@github-actions github-actions Bot Jul 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

important — rejections here are silent, and DeadlineExceeded is reported as "server busy"

Two things:

  1. No log on rejection. The HTTP path calls h.logServerBusy() (jsonrpc/http.go:60-66), which emits a sampled warning with running/queued/rejected. The websocket path drops the request with no log at all, so an operator debugging "clients randomly get -32603" has nothing to go on. Consider giving Websocket the same log.Sampled logger and calling it here.

  2. The two failure modes are conflated. Gate.Acquire returns either ErrServerBusy (queue full — a genuine overload) or context.DeadlineExceeded (this request queued longer than requestTimeout — the server is busy but this was a timeout). Both fall through to serverBusyResponse. HTTP distinguishes them (http.go:94-102). Worth at least an explicit errors.Is(err, ErrServerBusy) branch so the two are separable in logs/metrics rather than relying on "not context.Canceled" as the catch-all.

Also note the metrics side: makeHTTPGateMetrics (node/metrics.go:110-135) registers these under subsystemHTTP with help text "Number of HTTP RPC requests…". Now that the same gate serves websockets, rpc_http_active_requests / rpc_http_rejected_requests silently include websocket traffic. Either rename to a version-neutral subsystem or update the help strings, otherwise existing dashboards/alerts become wrong after this lands.

defer ws.gate.Release()
}

return ws.rpc.HandleReadWriter(wsc.ctx, ws.requestTimeout, wsc)
}
Comment on lines +167 to +186

@github-actions github-actions Bot Jul 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

important — a websocket request can now consume 2 × rpcRequestTimeout

acquireCtx gets its own requestTimeout deadline (line 172), and then HandleReadWriter derives a fresh requestTimeout from wsc.ctx (jsonrpc/server.go:322-323), not from acquireCtx. So a request that waits requestTimeout-ε in the gate queue still gets a full requestTimeout to execute — worst case double the configured budget, during which the connection's read loop (ServeHTTP line 132) is blocked and no other message on that connection makes progress.

The HTTP path doesn't have this: jsonrpc/http.go:80-105 builds one ctx with the timeout and uses it for both the gate acquire and the handler, so queue time is charged against the request budget.

Suggest building the deadline once here and threading it into HandleReadWriter, e.g. give HandleReadWriter the already-deadlined ctx and pass 0 for requestTimeout:

func (ws *Websocket) handleMessage(wsc *websocketConn) error {
	ctx := wsc.ctx
	if ws.requestTimeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, ws.requestTimeout)
		defer cancel()
	}
	if ws.gate != nil {
		if err := ws.gate.Acquire(ctx); err != nil { ... }
		defer ws.gate.Release()
	}
	return ws.rpc.HandleReadWriter(ctx, 0, wsc)
}

Careful with that shape though: HandleReadWriter passes connCtx into connection.ctx, which subscription writes use after the initial response, so handing it a deadlined ctx would break long-lived subscriptions. Either keep the two ctxs separate but subtract the queue wait, or pass the deadline explicitly. Worth a comment either way, since the current asymmetry with HTTP looks unintentional.


type WebsocketConnParams struct {
// Maximum message size allowed.
ReadLimit int64
Expand Down
56 changes: 56 additions & 0 deletions jsonrpc/websocket_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

@github-actions github-actions Bot Jul 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit — Running() == 0 is not sufficient to know the slot is fully reclaimed

Gate.Release() (jsonrpc/gate.go) does:

func (g *Gate) Release() {
	<-g.sem              // Running() drops to 0 here
	g.decreaseActiveReq() // activeRequests drops to 0 only here
}

Running() is len(g.sem), but admission is decided by activeRequests vs maxRequests. So there is a window where Running() == 0 while activeRequests is still 1. If Eventually samples inside that window, the follow-up id:3 request on connB hits increaseActiveReq() == 2 > maxRequests == 1 and gets serverBusyResponse instead of {"result":"hi"} — the test fails.

It's a nanosecond-wide window and there's a websocket round trip after it, so it will almost never fire in CI, but it's a real ordering gap rather than a theoretical one. Cheap fix — assert both counters are drained:

Suggested change
require.Eventually(t, func() bool { return gate.Running() == 0 }, time.Second, 5*time.Millisecond)
require.Eventually(t, func() bool { return gate.Running() == 0 && gate.Queued() == 0 },
time.Second, 5*time.Millisecond)

Fix this →


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))
}
18 changes: 4 additions & 14 deletions node/http.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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).
Expand Down Expand Up @@ -155,6 +143,7 @@ func makeRPCOverWebsocket(
metricsEnabled bool,
corsEnabled bool,
rpcRequestTimeout time.Duration,
gate *jsonrpc.Gate,
) *httpService {
var listener jsonrpc.NewRequestListener
if metricsEnabled {
Expand All @@ -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)
}
Expand Down
11 changes: 9 additions & 2 deletions node/node.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -543,6 +543,13 @@
"/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)

Check warning on line 550 in node/node.go

View check run for this annotation

Codecov / codecov/patch

node/node.go#L548-L550 

Added lines #L548 - L550 were not covered by tests
}
}
Comment on lines +546 to +552

@github-actions github-actions Bot Jul 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit — hoisting the gate here so HTTP and WS share it is the right call, but it now runs unconditionally: with --http=false --ws=false the node still builds a Gate and (with metrics on) registers three prometheus collectors for a gate nobody acquires. Harmless, but if cfg.HTTP || cfg.Websocket around it would keep /metrics honest.

if cfg.HTTP {
readinessHandlers := NewReadinessHandlers(chain, synchronizer, cfg.ReadinessBlockTolerance)
httpHandlers := map[string]http.HandlerFunc{
Expand All @@ -561,8 +568,7 @@
cfg.Metrics,
cfg.RPCCorsEnable,
cfg.RPCRequestTimeout,
cfg.RPCMaxConcurrentRequests,
cfg.RPCMaxRequestQueue,
rpcGate,
),
)
}
Expand All @@ -577,6 +583,7 @@
cfg.Metrics,
cfg.RPCCorsEnable,
cfg.RPCRequestTimeout,
rpcGate,
),
)
}
Expand Down
11 changes: 8 additions & 3 deletions rpc/handlers.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand All @@ -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)

@github-actions github-actions Bot Jul 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

important — the limit is process-wide, not per-connection

The PR description says "Bound the number of active subscriptions per connection", but this is a single semaphore.Weighted(2048) shared by all three handlers and, transitively, by every websocket connection in the process. There is no per-connection accounting anywhere in subscribe().

Combined with jsonrpc/websocket.go:21 (maxConns = 2048), a single client that opens one connection and issues 2048 starknet_subscribe* calls exhausts the global pool, and every other client — including already-connected ones — gets ErrTooManySubscriptions on their next subscribe. That's cheaper to trigger than the resource exhaustion this is meant to prevent.

Suggest tracking a per-connection counter (keyed on the jsonrpc.Conn, which subscribe already receives, and which has Equal for identity) and keeping the global cap as a second, larger backstop.

Fix this →

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,
Expand Down
3 changes: 3 additions & 0 deletions rpc/rpccore/rpccore.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@ const (
MaxBlocksBack = 1024
EntrypointNotFoundFelt string = "0x454e545259504f494e545f4e4f545f464f554e44"
ErrEPSNotFound = "Entry point EntryPointSelector(%s) not found in contract."

DefaultMaxSubscriptions int64 = 2048
Comment on lines +22 to +23

@github-actions github-actions Bot Jul 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit — hardcoded, unlike the sibling request gate

rpc-max-concurrent-requests / rpc-max-request-queue are both configurable flags (cmd/juno/juno.go:110,176), but the subscription cap is a compile-time constant with no way to tune or disable it. Node operators with different memory budgets can't adjust it, and there's no way to turn it off if it turns out to be too tight. Worth a rpc-max-subscriptions flag for symmetry — the plumbing already supports it since WithSubscriptionLimiter takes the semaphore.

Also 2048 deserves the same "arbitrary default, revisit after monitoring" note that maxConns carries in jsonrpc/websocket.go:21, and ErrTooManySubscriptions (code 101) is a new Juno-specific error that clients will hit — please add it to the RPC error docs.

@rodrodros rodrodros Jul 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Valid

@infrmtcs infrmtcs Jul 29, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In fact, if we make them configurable, it should be a pair of flags: 1 for connections and 1 for subscriptions.

)

//go:generate mockgen -destination=../mocks/mock_gateway_handler.go -package=mocks github.com/NethermindEth/juno/rpc/rpccore Gateway
Expand Down Expand Up @@ -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"}
)
11 changes: 9 additions & 2 deletions rpc/v10/handlers.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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?
Expand Down Expand Up @@ -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
Expand Down
6 changes: 6 additions & 0 deletions rpc/v10/subscriptions.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Comment on lines +57 to +59

@github-actions github-actions Bot Jul 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit — the limit is checked after the expensive work, and the pair is triplicated

The acquire/release pairing itself is correct: there is no early return between TryAcquire (line 57) and sub.wg.Go (line 80), the release is deferred first inside the goroutine so it runs after unsubscribe, and Unsubscribe does subs.wg.Wait() (line 230), so the slot is provably free before starknet_unsubscribe returns true. No leak, no double-release.

Two smaller things:

  1. Load is shed too late. All five Subscribe* entrypoints resolve the block range before calling subscribe, so a rejected subscription has already paid bcReader.Height() and — for historical subscriptions — the header/commitment/state-update reads. Your own test makes this visible: Height() is expected .Times(3) for two successful subscriptions. Checking the limiter at the top of the Subscribe* methods (or hoisting it into a tiny helper called first) would reject before touching the DB, which is the point of a limiter.

  2. Identical 6 lines now live in v8/subscriptions.go:121,138, v9/subscriptions.go:109,131 and here. The three subscribe bodies were already near-duplicates, so this follows precedent, but it means a future fix (e.g. the per-connection accounting discussed on rpc/handlers.go) has to be applied three times and can silently diverge. Worth a shared helper in rpccore if that's cheap here.

id := h.idgen()
//nolint:gosec // G118: cancel called in unsubscribe()
subscriptionCtx, subscriptionCtxCancel := context.WithCancel(wsConn.Context())
Expand All @@ -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)
Expand Down
44 changes: 44 additions & 0 deletions rpc/v10/subscriptions_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
}
Comment on lines +1221 to +1262

@github-actions github-actions Bot Jul 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit — good test, two gaps

Solid: it exercises reject-at-limit and slot-reclaimed-after-unsubscribe, and it correctly relies on Unsubscribe's wg.Wait() rather than a sleep, so it isn't flaky.

  • The identical limiter code in rpc/v8/subscriptions.go:121 and rpc/v9/subscriptions.go:109 has no coverage. Since it's copy-pasted three times, that's exactly where a divergence would go unnoticed.
  • Only SubscribeNewHeads is covered. SubscribeEvents / SubscribeTransactionStatus / SubscribeNewTransactions / SubscribeNewTransactionReceipts all funnel through the same subscribe, so a small table-driven test over the entrypoints would be cheap insurance that none of them grows an early return between TryAcquire and sub.wg.Go.

Also subID3/conn3 skips 2 — presumably left over from an earlier draft.


func TestSubscribeNewHeadsHistorical(t *testing.T) {
logger := log.NewNopZapLogger()
client := feeder.NewTestClient(t, &networks.Sepolia)
Expand Down
Loading
Loading