gito/services/core/internal/protosocket/server.go

package protosocket

import (
	"context"
	"log/slog"
	"net/http"
	"sync"
	"time"

	toki "git.toki-labs.com/toki/proto-socket/go"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/structpb"
	"nhooyr.io/websocket"
)

type Config struct {
	HeartbeatIntervalSec int
	HeartbeatWaitSec     int
}

type Server struct {
	cfg           Config
	dispatcher    *Dispatcher
	logger        *slog.Logger
	mu            sync.Mutex
	clients       map[*toki.WsClient]clientDiagnostics
	subscriptions map[string]EventSubscription
}

type clientDiagnostics struct {
	connectionID string
}

type clientSnapshot struct {
	client       *toki.WsClient
	diagnostic   clientDiagnostics
	subscription EventSubscription
	subscribed   bool
}

func NewServer(cfg Config, logger *slog.Logger) *Server {
	if cfg.HeartbeatIntervalSec <= 0 {
		cfg.HeartbeatIntervalSec = 30
	}
	if cfg.HeartbeatWaitSec <= 0 {
		cfg.HeartbeatWaitSec = 10
	}
	return &Server{
		cfg:           cfg,
		dispatcher:    NewDispatcher(),
		logger:        logger,
		clients:       make(map[*toki.WsClient]clientDiagnostics),
		subscriptions: make(map[string]EventSubscription),
	}
}

func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	opts := &websocket.AcceptOptions{InsecureSkipVerify: true}
	conn, err := websocket.Accept(w, r, opts)
	if err != nil {
		s.logError("failed to accept proto-socket connection", "error", err)
		return
	}

	connectionID := "conn-" + generateID()
	diagnostic := clientDiagnostics{connectionID: connectionID}
	client := toki.NewWsClient(
		conn,
		s.cfg.HeartbeatIntervalSec,
		s.cfg.HeartbeatWaitSec,
		ParserMap(),
	)

	s.mu.Lock()
	s.clients[client] = diagnostic
	s.mu.Unlock()

	client.AddDisconnectListener(func(c *toki.WsClient) {
		s.mu.Lock()
		diagnostic := s.clients[c]
		delete(s.clients, c)
		delete(s.subscriptions, diagnostic.connectionID)
		s.mu.Unlock()
		s.logInfo("proto-socket client disconnected", diagnosticsLogAttrs(Envelope{ProtocolVersion: ProtocolVersion}, diagnostic.connectionID)...)
	})

	toki.AddRequestListenerTyped[*structpb.Struct, *structpb.Struct](&client.Communicator, func(req *structpb.Struct) (*structpb.Struct, error) {
		env, err := EnvelopeFromStruct(req)
		if err != nil {
			errEnv := withDiagnosticsMeta(Envelope{
				ProtocolVersion: ProtocolVersion,
				ID:              generateID(),
				Type:            "error",
				Error: &EnvelopeError{
					Code:      "INVALID_ENVELOPE",
					Message:   "failed to parse envelope: " + err.Error(),
					Retryable: false,
				},
			}, connectionID)
			res, _ := errEnv.ToStruct()
			return res, nil
		}
		env = withConnectionMeta(env, connectionID)

		resEnv := withDiagnosticsMeta(s.dispatcher.Dispatch(r.Context(), env), connectionID)
		resStruct, err := resEnv.ToStruct()
		if err != nil {
			errEnv := withDiagnosticsMeta(ErrorResponse(env, "SERIALIZATION_ERROR", "failed to serialize response envelope: "+err.Error(), false), connectionID)
			res, _ := errEnv.ToStruct()
			s.logError("proto-socket response serialization failed", append(diagnosticsLogAttrs(errEnv, connectionID), "error", err)...)
			return res, nil
		}

		s.logInfo("proto-socket request handled", diagnosticsLogAttrs(resEnv, connectionID)...)
		return resStruct, nil
	})

	s.logInfo("proto-socket client connected", diagnosticsLogAttrs(Envelope{ProtocolVersion: ProtocolVersion}, connectionID)...)
}

func (s *Server) BroadcastEnvelope(ctx context.Context, env Envelope) error {
	s.mu.Lock()
	clients := make([]clientSnapshot, 0, len(s.clients))
	for client, diagnostic := range s.clients {
		subscription, subscribed := s.subscriptions[diagnostic.connectionID]
		clients = append(clients, clientSnapshot{
			client:       client,
			diagnostic:   diagnostic,
			subscription: subscription,
			subscribed:   subscribed,
		})
	}
	s.mu.Unlock()

	for _, snapshot := range clients {
		if !snapshot.client.IsAlive() {
			continue
		}
		if !snapshot.subscribed || !snapshot.subscription.Allows(env) {
			continue
		}
		eventEnv := withDiagnosticsMeta(env, snapshot.diagnostic.connectionID)
		msg, err := eventEnv.ToStruct()
		if err != nil {
			s.logError("proto-socket broadcast serialization failed", append(diagnosticsLogAttrs(eventEnv, snapshot.diagnostic.connectionID), "error", err)...)
			return err
		}
		if err := snapshot.client.Send(msg); err != nil {
			s.logError("failed to send proto-socket broadcast", append(diagnosticsLogAttrs(eventEnv, snapshot.diagnostic.connectionID), "error", err)...)
			continue
		}
		s.logInfo("proto-socket broadcast sent", diagnosticsLogAttrs(eventEnv, snapshot.diagnostic.connectionID)...)
	}
	return nil
}

func (s *Server) Subscribe(connectionID string, subscription EventSubscription) EventSubscription {
	subscription = subscription.Normalized()
	s.mu.Lock()
	defer s.mu.Unlock()
	s.subscriptions[connectionID] = subscription
	return subscription
}

func (s *Server) Dispatcher() *Dispatcher {
	return s.dispatcher
}

func ParserMap() toki.ParserMap {
	return toki.ParserMap{
		toki.TypeNameOf(&structpb.Struct{}): func(b []byte) (proto.Message, error) {
			msg := &structpb.Struct{}
			return msg, proto.Unmarshal(b, msg)
		},
	}
}

func withDiagnosticsMeta(env Envelope, connectionID string) Envelope {
	if env.ProtocolVersion == "" {
		env.ProtocolVersion = ProtocolVersion
	}
	meta := make(map[string]any, len(env.Meta)+6)
	for k, v := range env.Meta {
		meta[k] = v
	}
	meta["connection_id"] = connectionID
	meta["protocol_version"] = env.ProtocolVersion
	meta["channel"] = env.Channel
	meta["action"] = env.Action
	meta["error_code"] = errorCode(env)
	meta["timestamp"] = diagnosticsTimestamp()
	env.Meta = meta
	return env
}

func withConnectionMeta(env Envelope, connectionID string) Envelope {
	if env.ProtocolVersion == "" {
		env.ProtocolVersion = ProtocolVersion
	}
	meta := make(map[string]any, len(env.Meta)+2)
	for k, v := range env.Meta {
		meta[k] = v
	}
	meta["connection_id"] = connectionID
	meta["protocol_version"] = env.ProtocolVersion
	env.Meta = meta
	return env
}

func errorCode(env Envelope) string {
	if env.Error == nil {
		return ""
	}
	return env.Error.Code
}

func diagnosticsLogAttrs(env Envelope, connectionID string) []any {
	protocolVersion := env.ProtocolVersion
	if protocolVersion == "" {
		protocolVersion = ProtocolVersion
	}
	return []any{
		"connection_id", connectionID,
		"protocol_version", protocolVersion,
		"channel", env.Channel,
		"action", env.Action,
		"error_code", errorCode(env),
		"timestamp", diagnosticsTimestamp(),
	}
}

func diagnosticsTimestamp() string {
	return time.Now().UTC().Format(time.RFC3339Nano)
}

func (s *Server) logInfo(msg string, args ...any) {
	if s.logger != nil {
		s.logger.Info(msg, args...)
	}
}

func (s *Server) logError(msg string, args ...any) {
	if s.logger != nil {
		s.logger.Error(msg, args...)
	}
}