prysm/validator/helpers/node_connection.go

package helpers

import (
	"time"

	"google.golang.org/grpc"
)

// Use an interface with a private dummy function to force all other packages to call NewNodeConnection
type NodeConnection interface {
	GetGrpcClientConn() *grpc.ClientConn
	GetBeaconApiUrl() string
	GetBeaconApiHeaders() map[string][]string
	setBeaconApiHeaders(map[string][]string)
	GetBeaconApiTimeout() time.Duration
	setBeaconApiTimeout(time.Duration)
	// GetGrpcConnectionProvider returns the gRPC connection provider for multi-endpoint support.
	// Returns nil if no provider is configured (single endpoint mode).
	GetGrpcConnectionProvider() GrpcConnectionProvider
	setGrpcConnectionProvider(GrpcConnectionProvider)
	dummy()
}

type nodeConnection struct {
	grpcClientConn         *grpc.ClientConn
	grpcConnectionProvider GrpcConnectionProvider
	beaconApiUrl           string
	beaconApiHeaders       map[string][]string
	beaconApiTimeout       time.Duration
}

// NodeConnectionOption is a functional option for configuring the node connection.
type NodeConnectionOption func(nc NodeConnection)

// WithBeaconApiHeaders sets the HTTP headers that should be sent to the server along with each request.
func WithBeaconApiHeaders(headers map[string][]string) NodeConnectionOption {
	return func(nc NodeConnection) {
		nc.setBeaconApiHeaders(headers)
	}
}

// WithBeaconApiTimeout sets the HTTP request timeout.
func WithBeaconApiTimeout(timeout time.Duration) NodeConnectionOption {
	return func(nc NodeConnection) {
		nc.setBeaconApiTimeout(timeout)
	}
}

// WithGrpcConnectionProvider sets the gRPC connection provider for multi-endpoint support.
func WithGrpcConnectionProvider(provider GrpcConnectionProvider) NodeConnectionOption {
	return func(nc NodeConnection) {
		nc.setGrpcConnectionProvider(provider)
	}
}

func (c *nodeConnection) GetGrpcClientConn() *grpc.ClientConn {
	// If a connection provider is configured, use its current connection
	if c.grpcConnectionProvider != nil {
		return c.grpcConnectionProvider.CurrentConn()
	}
	return c.grpcClientConn
}

func (c *nodeConnection) GetBeaconApiUrl() string {
	return c.beaconApiUrl
}

func (c *nodeConnection) GetBeaconApiHeaders() map[string][]string {
	return c.beaconApiHeaders
}

func (c *nodeConnection) setBeaconApiHeaders(headers map[string][]string) {
	c.beaconApiHeaders = headers
}

func (c *nodeConnection) GetBeaconApiTimeout() time.Duration {
	return c.beaconApiTimeout
}

func (c *nodeConnection) setBeaconApiTimeout(timeout time.Duration) {
	c.beaconApiTimeout = timeout
}

func (c *nodeConnection) GetGrpcConnectionProvider() GrpcConnectionProvider {
	return c.grpcConnectionProvider
}

func (c *nodeConnection) setGrpcConnectionProvider(provider GrpcConnectionProvider) {
	c.grpcConnectionProvider = provider
}

func (*nodeConnection) dummy() {}

func NewNodeConnection(grpcConn *grpc.ClientConn, beaconApiUrl string, opts ...NodeConnectionOption) NodeConnection {
	conn := &nodeConnection{}
	conn.grpcClientConn = grpcConn
	conn.beaconApiUrl = beaconApiUrl
	for _, opt := range opts {
		opt(conn)
	}
	return conn
}