[Access] Create draft ws brocker for new PUB/SUB (WIP/PoC)

cmd/access/node_builder/access_node_builder.go

@@ -146,6 +146,7 @@ type AccessNodeConfig struct {
 	apiBurstlimits                       map[string]int
 	rpcConf                              rpc.Config
 	stateStreamConf                      statestreambackend.Config
+	wsConfig                             routes.WebsocketConfig
 	stateStreamFilterConf                map[string]int
 	ExecutionNodeAddress                 string // deprecated
 	HistoricalAccessRPCs                 []access.AccessAPIClient
@@ -238,7 +239,12 @@ func DefaultAccessNodeConfig() *AccessNodeConfig {
 			ResponseLimit:           subscription.DefaultResponseLimit,
 			HeartbeatInterval:       subscription.DefaultHeartbeatInterval,
 		},
-		stateStreamFilterConf:        nil,
+		stateStreamFilterConf: nil,
+		wsConfig: routes.WebsocketConfig{
+			MaxSubscriptionsPerConnection: routes.DefaultMaxSubscriptionsPerConnection,
+			MaxResponsesPerSecond:         routes.DefaultMaxResponsesPerSecond,
+			SendMessageTimeout:            routes.DefaultSendMessageTimeout,
+		},
 		ExecutionNodeAddress:         "localhost:9000",
 		logTxTimeToFinalized:         false,
 		logTxTimeToExecuted:          false,
@@ -1436,6 +1442,23 @@ func (builder *FlowAccessNodeBuilder) extraFlags() {
 			"registerdb-pruning-threshold",
 			defaultConfig.registerDBPruneThreshold,
 			fmt.Sprintf("specifies the number of blocks below the latest stored block height to keep in register db. default: %d", defaultConfig.registerDBPruneThreshold))
+
+		// Websocket subscriptions
+		flags.Uint64Var(&builder.wsConfig.MaxSubscriptionsPerConnection,
+			"websocket-max-subscriptions-per-connection",
+			defaultConfig.wsConfig.MaxSubscriptionsPerConnection,
+			fmt.Sprintf("maximum number of subscriptions per connection for websocket subscriptions. Default: %d", builder.wsConfig.MaxSubscriptionsPerConnection))
+
+		flags.Uint64Var(&builder.wsConfig.MaxResponsesPerSecond,
+			"websocket-max-responses-per-second",
+			defaultConfig.wsConfig.MaxResponsesPerSecond,
+			fmt.Sprintf("maximum number of responses per second for websocket subscriptions. Default: %d", builder.wsConfig.MaxResponsesPerSecond))
+
+		flags.DurationVar(&builder.wsConfig.SendMessageTimeout,
+			"websocket-send-message-timeout",
+			defaultConfig.wsConfig.SendMessageTimeout,
+			fmt.Sprintf("timeout value for send messages for websocket subscriptions. Default: %v", defaultConfig.wsConfig.SendMessageTimeout))
+
 	}).ValidateFlags(func() error {
 		if builder.supportsObserver && (builder.PublicNetworkConfig.BindAddress == cmd.NotSet || builder.PublicNetworkConfig.BindAddress == "") {
 			return errors.New("public-network-address must be set if supports-observer is true")
@@ -2015,6 +2038,7 @@ func (builder *FlowAccessNodeBuilder) Build() (cmd.Node, error) {
 				builder.unsecureGrpcServer,
 				builder.stateStreamBackend,
 				builder.stateStreamConf,
+				builder.wsConfig,
 				indexReporter,
 			)
 			if err != nil {

cmd/observer/node_builder/observer_builder.go

@@ -139,6 +139,7 @@ type ObserverServiceConfig struct {
 	checkpointFile                       string
 	apiTimeout                           time.Duration
 	stateStreamConf                      statestreambackend.Config
+	wsConfig                             routes.WebsocketConfig
 	stateStreamFilterConf                map[string]int
 	upstreamNodeAddresses                []string
 	upstreamNodePublicKeys               []string
@@ -208,7 +209,12 @@ func DefaultObserverServiceConfig() *ObserverServiceConfig {
 			HeartbeatInterval:       subscription.DefaultHeartbeatInterval,
 			RegisterIDsRequestLimit: state_stream.DefaultRegisterIDsRequestLimit,
 		},
-		stateStreamFilterConf:                nil,
+		stateStreamFilterConf: nil,
+		wsConfig: routes.WebsocketConfig{
+			MaxSubscriptionsPerConnection: routes.DefaultMaxSubscriptionsPerConnection,
+			MaxResponsesPerSecond:         routes.DefaultMaxResponsesPerSecond,
+			SendMessageTimeout:            routes.DefaultSendMessageTimeout,
+		},
 		rpcMetricsEnabled:                    false,
 		apiRatelimits:                        nil,
 		apiBurstlimits:                       nil,
@@ -798,6 +804,23 @@ func (builder *ObserverServiceBuilder) extraFlags() {
 			"registerdb-pruning-threshold",
 			defaultConfig.registerDBPruneThreshold,
 			fmt.Sprintf("specifies the number of blocks below the latest stored block height to keep in register db. default: %d", defaultConfig.registerDBPruneThreshold))
+
+		// Websocket subscriptions
+		flags.Uint64Var(&builder.wsConfig.MaxSubscriptionsPerConnection,
+			"websocket-max-subscriptions-per-connection",
+			defaultConfig.wsConfig.MaxSubscriptionsPerConnection,
+			fmt.Sprintf("maximum number of subscriptions per connection for websocket subscriptions. Default: %d", builder.wsConfig.MaxSubscriptionsPerConnection))
+
+		flags.Uint64Var(&builder.wsConfig.MaxResponsesPerSecond,
+			"websocket-max-responses-per-second",
+			defaultConfig.wsConfig.MaxResponsesPerSecond,
+			fmt.Sprintf("maximum number of responses per second for websocket subscriptions. Default: %d", builder.wsConfig.MaxResponsesPerSecond))
+
+		flags.DurationVar(&builder.wsConfig.SendMessageTimeout,
+			"websocket-send-message-timeout",
+			defaultConfig.wsConfig.SendMessageTimeout,
+			fmt.Sprintf("timeout value for send messages for websocket subscriptions. Default: %v", defaultConfig.wsConfig.SendMessageTimeout))
+
 	}).ValidateFlags(func() error {
 		if builder.executionDataSyncEnabled {
 			if builder.executionDataConfig.FetchTimeout <= 0 {
@@ -1931,6 +1954,7 @@ func (builder *ObserverServiceBuilder) enqueueRPCServer() {
 			builder.unsecureGrpcServer,
 			builder.stateStreamBackend,
 			builder.stateStreamConf,
+			builder.wsConfig,
 			indexReporter,
 		)
 		if err != nil {

cmd/util/cmd/run-script/cmd.go

@@ -16,6 +16,7 @@ import (
 	"github.com/onflow/flow-go/cmd/util/ledger/util"
 	"github.com/onflow/flow-go/cmd/util/ledger/util/registers"
 	"github.com/onflow/flow-go/engine/access/rest"
+	"github.com/onflow/flow-go/engine/access/rest/routes"
 	"github.com/onflow/flow-go/engine/access/state_stream/backend"
 	"github.com/onflow/flow-go/engine/access/subscription"
 	"github.com/onflow/flow-go/engine/execution/computation"
@@ -169,6 +170,7 @@ func run(*cobra.Command, []string) {
 			metrics.NewNoopCollector(),
 			nil,
 			backend.Config{},
+			routes.WebsocketConfig{},
 		)
 		if err != nil {
 			log.Fatal().Err(err).Msg("failed to create server")

engine/access/handle_irrecoverable_state_test.go

@@ -171,6 +171,7 @@ func (suite *IrrecoverableStateTestSuite) SetupTest() {
 		suite.unsecureGrpcServer,
 		nil,
 		stateStreamConfig,
+		routes.WebsocketConfig{},
 		nil,
 	)
 	assert.NoError(suite.T(), err)

engine/access/integration_unsecure_grpc_server_test.go

@@ -21,6 +21,7 @@ import (
 	"github.com/onflow/flow-go/engine"
 	"github.com/onflow/flow-go/engine/access/index"
 	accessmock "github.com/onflow/flow-go/engine/access/mock"
+	"github.com/onflow/flow-go/engine/access/rest/routes"
 	"github.com/onflow/flow-go/engine/access/rpc"
 	"github.com/onflow/flow-go/engine/access/rpc/backend"
 	"github.com/onflow/flow-go/engine/access/state_stream"
@@ -215,6 +216,7 @@ func (suite *SameGRPCPortTestSuite) SetupTest() {
 		suite.unsecureGrpcServer,
 		nil,
 		stateStreamConfig,
+		routes.WebsocketConfig{},
 		nil,
 	)
 	assert.NoError(suite.T(), err)

engine/access/rest/routes/router.go

@@ -12,6 +12,7 @@ import (
 	"github.com/onflow/flow-go/access"
 	"github.com/onflow/flow-go/engine/access/rest/middleware"
 	"github.com/onflow/flow-go/engine/access/rest/models"
+	"github.com/onflow/flow-go/engine/access/rest/routes/subscription_handlers"
 	"github.com/onflow/flow-go/engine/access/state_stream"
 	"github.com/onflow/flow-go/engine/access/state_stream/backend"
 	"github.com/onflow/flow-go/model/flow"
@@ -78,6 +79,21 @@ func (b *RouterBuilder) AddWsRoutes(
 	return b
 }
 
+// AddPubSubRoute adds WebSocket route for the pub/sub mechanism to the router.
+func (b *RouterBuilder) AddPubSubRoute(
+	chain flow.Chain,
+	wsConfig WebsocketConfig,
+	subHandlerFactory *subscription_handlers.SubscriptionHandlerFactory,
+) *RouterBuilder {
+	b.v1SubRouter.
+		Methods(http.MethodGet).
+		Path("/ws").
+		Name("ws").
+		Handler(NewWSBrokerHandler(b.logger, wsConfig, chain, subHandlerFactory))
+
+	return b
+}
+
 func (b *RouterBuilder) Build() *mux.Router {
 	return b.router
 }

engine/access/rest/routes/subscription_handlers/subscription_handler.go

@@ -0,0 +1,53 @@
+package subscription_handlers
+
+import (
+	"fmt"
+
+	"github.com/onflow/flow-go/access"
+	"github.com/onflow/flow-go/engine/access/state_stream"
+)
+
+const (
+	EventsTopic              = "events"
+	AccountStatusesTopic     = "account_statuses"
+	BlocksTopic              = "blocks"
+	BlockHeadersTopic        = "block_headers"
+	BlockDigestsTopic        = "block_digests"
+	TransactionStatusesTopic = "transaction_statuses"
+)
+
+type SubscriptionHandler interface {
+	ID() string
+	Topic() string
+	Close() error
+}
+
+type SubscriptionHandlerFactory struct {
+	eventFilterConfig state_stream.EventFilterConfig
+
+	stateStreamApi state_stream.API
+	accessApi      access.API
+}
+
+func NewSubscriptionHandlerFactory(eventFilterConfig state_stream.EventFilterConfig, stateStreamApi state_stream.API, accessApi access.API) *SubscriptionHandlerFactory {
+	return &SubscriptionHandlerFactory{
+		eventFilterConfig: eventFilterConfig,
+		stateStreamApi:    stateStreamApi,
+		accessApi:         accessApi,
+	}
+}
+
+func (s *SubscriptionHandlerFactory) CreateSubscriptionHandler(topic string, arguments map[string]interface{}, broadcastMessage func(interface{})) (SubscriptionHandler, error) {
+	switch topic {
+	// TODO: Implemented handlers for each topic should be added in respective case
+	case EventsTopic,
+		AccountStatusesTopic,
+		BlocksTopic,
+		BlockHeadersTopic,
+		BlockDigestsTopic,
+		TransactionStatusesTopic:
+		return nil, fmt.Errorf("topic \"%s\" not implemented yet", topic)
+	default:
+		return nil, fmt.Errorf("unsupported topic \"%s\"", topic)
+	}
+}