Enforce global stream limit per WebSocket subscription

cmd/access/node_builder/access_node_builder.go

@@ -1763,6 +1763,10 @@ func (builder *FlowAccessNodeBuilder) extraFlags() {
 			return errors.New("execution-data-indexing-enabled must be set if store-tx-result-error-messages is enabled")
 		}
 
+		if builder.stateStreamConf.MaxGlobalStreams == 0 {
+			return errors.New("state-stream-global-max-streams must be greater than 0")
+		}
+
 		return nil
 	})
 }

cmd/observer/node_builder/observer_builder.go

@@ -945,6 +945,10 @@ func (builder *ObserverServiceBuilder) extraFlags() {
 			return errors.New("rest-max-request-size must be greater than 0")
 		}
 
+		if builder.stateStreamConf.MaxGlobalStreams == 0 {
+			return errors.New("state-stream-global-max-streams must be greater than 0")
+		}
+
 		return nil
 	})
 }

engine/access/rest/router/router.go

@@ -108,15 +108,14 @@ func (b *RouterBuilder) AddWebsocketsRoute(
 	maxRequestSize int64,
 	maxResponseSize int64,
 	dataProviderFactory dp.DataProviderFactory,
-	limiter *limiters.ConcurrencyLimiter,
+	streamLimiter *limiters.ConcurrencyLimiter,
 ) *RouterBuilder {
-	h := websockets.NewWebSocketHandler(ctx, b.logger, config, chain, maxRequestSize, maxResponseSize, dataProviderFactory)
-	handler := websockets.NewConnectionLimitedHandler(b.logger, h.HttpHandler, h, limiter)
+	h := websockets.NewWebSocketHandler(ctx, b.logger, config, chain, maxRequestSize, maxResponseSize, dataProviderFactory, streamLimiter)
 	b.v1SubRouter.
 		Methods(http.MethodGet).
 		Path("/ws").
 		Name("ws").
-		Handler(handler)
+		Handler(h)
 
 	return b
 }

engine/access/rest/server.go

@@ -1,6 +1,7 @@
 package rest
 
 import (
+	"errors"
 	"net/http"
 	"time"
 
@@ -55,6 +56,10 @@ func NewServer(
 	extendedBackend extended.API,
 	limiter *limiters.ConcurrencyLimiter,
 ) (*http.Server, error) {
+	if limiter == nil && (stateStreamApi != nil || enableNewWebsocketsStreamAPI) {
+		return nil, errors.New("stream limiter is required when websocket routes are enabled")
+	}
+
 	builder := router.NewRouterBuilder(logger, restCollector).AddRestRoutes(serverAPI, chain, config.MaxRequestSize, config.MaxResponseSize)
 	if stateStreamApi != nil {
 		builder.AddLegacyWebsocketsRoutes(stateStreamApi, chain, stateStreamConfig, config.MaxRequestSize, config.MaxResponseSize, limiter)

engine/access/rest/websockets/controller.go

@@ -89,6 +89,7 @@ import (
 
 	dp "github.com/onflow/flow-go/engine/access/rest/websockets/data_providers"
 	"github.com/onflow/flow-go/engine/access/rest/websockets/models"
+	"github.com/onflow/flow-go/module/limiters"
 	"github.com/onflow/flow-go/utils/concurrentmap"
 )
 
@@ -135,7 +136,8 @@ type Controller struct {
 	dataProviders       *concurrentmap.Map[SubscriptionID, dp.DataProvider]
 	dataProviderFactory dp.DataProviderFactory
 	dataProvidersGroup  *sync.WaitGroup
-	limiter             *rate.Limiter
+	rateLimiter         *rate.Limiter
+	streamLimiter       *limiters.ConcurrencyLimiter
 
 	keepaliveConfig KeepaliveConfig
 }
@@ -145,10 +147,15 @@ func NewWebSocketController(
 	config Config,
 	conn WebsocketConnection,
 	dataProviderFactory dp.DataProviderFactory,
-) *Controller {
-	var limiter *rate.Limiter
+	streamLimiter *limiters.ConcurrencyLimiter,
+) (*Controller, error) {
+	if streamLimiter == nil {
+		return nil, errors.New("stream limiter is required")
+	}
+
+	var rateLimiter *rate.Limiter
 	if config.MaxResponsesPerSecond > 0 {
-		limiter = rate.NewLimiter(rate.Limit(config.MaxResponsesPerSecond), 1)
+		rateLimiter = rate.NewLimiter(rate.Limit(config.MaxResponsesPerSecond), 1)
 	}
 
 	return &Controller{
@@ -159,9 +166,10 @@ func NewWebSocketController(
 		dataProviders:       concurrentmap.New[SubscriptionID, dp.DataProvider](),
 		dataProviderFactory: dataProviderFactory,
 		dataProvidersGroup:  &sync.WaitGroup{},
-		limiter:             limiter,
+		rateLimiter:         rateLimiter,
+		streamLimiter:       streamLimiter,
 		keepaliveConfig:     DefaultKeepaliveConfig(),
-	}
+	}, nil
 }
 
 // HandleConnection manages the lifecycle of a WebSocket connection,
@@ -442,8 +450,20 @@ func (c *Controller) handleSubscribe(ctx context.Context, msg models.SubscribeMe
 		return
 	}
 
+	// Check if the global stream limit has been reached.
+	if !c.streamLimiter.Acquire() {
+		err := fmt.Errorf("error creating new subscription: maximum number of streams reached")
+		c.writeErrorResponse(
+			ctx,
+			err,
+			wrapErrorMessage(http.StatusTooManyRequests, err.Error(), models.SubscribeAction, msg.SubscriptionID),
+		)
+		return
+	}
+
 	subscriptionID, err := c.parseOrCreateSubscriptionID(msg.SubscriptionID)
 	if err != nil {
+		c.streamLimiter.Release()
 		err = fmt.Errorf("error parsing subscription id: %w", err)
 		c.writeErrorResponse(
 			ctx,
@@ -456,6 +476,7 @@ func (c *Controller) handleSubscribe(ctx context.Context, msg models.SubscribeMe
 	// register new provider
 	provider, err := c.dataProviderFactory.NewDataProvider(ctx, subscriptionID.String(), msg.Topic, msg.Arguments, c.multiplexedStream)
 	if err != nil {
+		c.streamLimiter.Release()
 		err = fmt.Errorf("error creating data provider: %w", err)
 		c.writeErrorResponse(
 			ctx,
@@ -478,6 +499,8 @@ func (c *Controller) handleSubscribe(ctx context.Context, msg models.SubscribeMe
 	// run provider
 	c.dataProvidersGroup.Add(1)
 	go func() {
+		defer c.streamLimiter.Release()
+
 		err = provider.Run()
 		if err != nil {
 			err = fmt.Errorf("internal error: %w", err)
@@ -604,9 +627,9 @@ func (c *Controller) parseOrCreateSubscriptionID(id string) (SubscriptionID, err
 // An error is returned if the context is canceled or the expected wait time exceeds the context's
 // deadline.
 func (c *Controller) checkRateLimit(ctx context.Context) error {
-	if c.limiter == nil {
+	if c.rateLimiter == nil {
 		return nil
 	}
 
-	return c.limiter.WaitN(ctx, 1)
+	return c.rateLimiter.WaitN(ctx, 1)
 }