Refactor data service: rename streaming files, update gRPC service methods, and enhance stream management

services/data_service/cmd/main.go

@@ -1,6 +1,10 @@
 package main
 
 import (
+	"fmt"
+	"log"
+	"net"
+
 	pb "gitlab.michelsen.id/phillmichelsen/tessera/pkg/pb/data_service"
 	"gitlab.michelsen.id/phillmichelsen/tessera/services/data_service/internal/manager"
 	"gitlab.michelsen.id/phillmichelsen/tessera/services/data_service/internal/provider/binance"
@@ -8,41 +12,55 @@ import (
 	"gitlab.michelsen.id/phillmichelsen/tessera/services/data_service/internal/server"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/reflection"
-	"log"
-	"net"
 )
 
 func main() {
+	fmt.Println("Starting Data Service...")
 	// Setup
 	r := router.NewRouter()
 	m := manager.NewManager(r)
 	binanceFutures := binance.NewFuturesWebsocket()
 	m.AddProvider("binance_futures_websocket", binanceFutures)
 
-	// gRPC Server
-	grpcServer := grpc.NewServer()
+	// gRPC Control Server
+	grpcControlServer := grpc.NewServer()
 	go func() {
-		pb.RegisterDataServiceStreamingServer(grpcServer, server.NewGRPCStreamingServer(m))
-		reflection.Register(grpcServer)
+		pb.RegisterDataServiceControlServer(grpcControlServer, server.NewGRPCControlServer(m))
+		reflection.Register(grpcControlServer)
 		grpcLis, err := net.Listen("tcp", ":50051")
+		if err != nil {
+			log.Fatalf("Failed to listen for gRPC control: %v", err)
+		}
+		log.Println("gRPC control server listening on :50051")
+		if err := grpcControlServer.Serve(grpcLis); err != nil {
+			log.Fatalf("Failed to serve gRPC control: %v", err)
+		}
+	}()
+
+	// gRPC Streaming Server
+	grpcStreamingServer := grpc.NewServer()
+	go func() {
+		pb.RegisterDataServiceStreamingServer(grpcStreamingServer, server.NewGRPCStreamingServer(m))
+		reflection.Register(grpcStreamingServer)
+		grpcLis, err := net.Listen("tcp", ":50052")
 		if err != nil {
 			log.Fatalf("Failed to listen for gRPC: %v", err)
 		}
-		log.Println("gRPC server listening on :50051")
-		if err := grpcServer.Serve(grpcLis); err != nil {
+		log.Println("gRPC streaming server listening on :50052")
+		if err := grpcStreamingServer.Serve(grpcLis); err != nil {
 			log.Fatalf("Failed to serve gRPC: %v", err)
 		}
 	}()
 
-	// Socket Server
-	socketServer := server.NewSocketStreamingServer(m)
+	// Socket Streaming Server
+	socketStreamingServer := server.NewSocketStreamingServer(m)
 	go func() {
 		socketLis, err := net.Listen("tcp", ":6000")
 		if err != nil {
 			log.Fatalf("Failed to listen for socket: %v", err)
 		}
 		log.Println("Socket server listening on :6000")
-		if err := socketServer.Serve(socketLis); err != nil {
+		if err := socketStreamingServer.Serve(socketLis); err != nil {
 			log.Fatalf("Socket server error: %v", err)
 		}
 	}()

services/data_service/internal/manager/manager.go

@@ -2,12 +2,13 @@ package manager
 
 import (
 	"fmt"
+	"sync"
+	"time"
+
 	"github.com/google/uuid"
 	"gitlab.michelsen.id/phillmichelsen/tessera/services/data_service/internal/domain"
 	"gitlab.michelsen.id/phillmichelsen/tessera/services/data_service/internal/provider"
 	"gitlab.michelsen.id/phillmichelsen/tessera/services/data_service/internal/router"
-	"sync"
-	"time"
 )
 
 type Manager struct {
@@ -38,59 +39,118 @@ func NewManager(router *router.Router) *Manager {
 	}
 }
 
-func (m *Manager) StartStream(ids []domain.Identifier) (uuid.UUID, error) {
+func (m *Manager) StartStream() (uuid.UUID, error) {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
-	// Validate Identifiers, prevents unnecessary calls to providers
-	for _, id := range ids {
-		if id.Provider == "" || id.Subject == "" {
-			return uuid.Nil, fmt.Errorf("invalid identifier: %v", id)
-		}
-		if _, ok := m.providers[id.Provider]; !ok {
-			return uuid.Nil, fmt.Errorf("unknown provider: %s", id.Provider)
-		}
-		if !m.providers[id.Provider].IsValidSubject(id.Subject, false) {
-			return uuid.Nil, fmt.Errorf("invalid subject for provider %s: %s", id.Provider, id.Subject)
-		}
-	}
-
-	// Provision provider streams
-	for _, id := range ids {
-		if _, ok := m.providerStreams[id]; ok {
-			continue // Skip if requested stream is already being provided
-		}
-
-		ch := make(chan domain.Message, 64)
-		if err := m.providers[id.Provider].RequestStream(id.Subject, ch); err != nil {
-			return uuid.Nil, fmt.Errorf("provision %v: %w", id, err)
-		}
-		m.providerStreams[id] = ch
-
-		// Start routine to route the provider stream to the router's input channel
-		go func(ch chan domain.Message) {
-			for msg := range ch {
-				m.router.IncomingChannel() <- msg
-			}
-		}(ch)
-	}
-
 	streamID := uuid.New()
-
 	m.clientStreams[streamID] = &ClientStream{
 		UUID:        streamID,
-		Identifiers: ids,
-		OutChannel:  nil, // Initially nil, will be set when connected
+		Identifiers: nil, // start empty
+		OutChannel:  nil, // not yet connected
 		Timer: time.AfterFunc(1*time.Minute, func() {
 			fmt.Printf("stream %s expired due to inactivity\n", streamID)
-			m.StopStream(streamID)
+			err := m.StopStream(streamID)
+			if err != nil {
+				fmt.Printf("failed to stop stream after timeout: %v\n", err)
+			}
 		}),
 	}
 
 	return streamID, nil
 }
 
-func (m *Manager) StopStream(streamID uuid.UUID) {
+func (m *Manager) ConfigureStream(streamID uuid.UUID, newIds []domain.Identifier) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	stream, ok := m.clientStreams[streamID]
+	if !ok {
+		return fmt.Errorf("stream not found: %s", streamID)
+	}
+
+	// Validate new identifiers.
+	for _, id := range newIds {
+		if id.Provider == "" || id.Subject == "" {
+			return fmt.Errorf("empty identifier: %v", id)
+		}
+		prov, exists := m.providers[id.Provider]
+		if !exists {
+			return fmt.Errorf("unknown provider: %s", id.Provider)
+		}
+		if !prov.IsValidSubject(id.Subject, false) {
+			return fmt.Errorf("invalid subject %q for provider %s", id.Subject, id.Provider)
+		}
+	}
+
+	// Generate old and new sets of identifiers
+	oldSet := make(map[domain.Identifier]struct{}, len(stream.Identifiers))
+	for _, id := range stream.Identifiers {
+		oldSet[id] = struct{}{}
+	}
+	newSet := make(map[domain.Identifier]struct{}, len(newIds))
+	for _, id := range newIds {
+		newSet[id] = struct{}{}
+	}
+
+	// Add identifiers that are in newIds but not in oldSet
+	for _, id := range newIds {
+		if _, seen := oldSet[id]; !seen {
+			// Provision the stream from the provider if needed
+			if _, ok := m.providerStreams[id]; !ok {
+				ch := make(chan domain.Message, 64)
+				if err := m.providers[id.Provider].RequestStream(id.Subject, ch); err != nil {
+					return fmt.Errorf("provision %v: %w", id, err)
+				}
+				m.providerStreams[id] = ch
+
+				incomingChannel := m.router.IncomingChannel()
+				go func(c chan domain.Message) {
+					for msg := range c {
+						incomingChannel <- msg
+					}
+				}(ch)
+			}
+
+			// Register the new identifier with the router, only if there's an active output channel (meaning the stream is connected)
+			if stream.OutChannel != nil {
+				m.router.RegisterRoute(id, stream.OutChannel)
+			}
+		}
+	}
+
+	// Remove identifiers that are in oldSet but not in newSet
+	for _, oldId := range stream.Identifiers {
+		if _, keep := newSet[oldId]; !keep {
+			// Deregister the identifier from the router, only if there's an active output channel (meaning the stream is connected)
+			if stream.OutChannel != nil {
+				m.router.DeregisterRoute(oldId, stream.OutChannel)
+			}
+		}
+	}
+
+	// Set the new identifiers for the stream
+	stream.Identifiers = newIds
+
+	// Clean up provider streams that are no longer used
+	used := make(map[domain.Identifier]bool)
+	for _, cs := range m.clientStreams {
+		for _, id := range cs.Identifiers {
+			used[id] = true
+		}
+	}
+	for id, ch := range m.providerStreams {
+		if !used[id] {
+			m.providers[id.Provider].CancelStream(id.Subject)
+			close(ch)
+			delete(m.providerStreams, id)
+		}
+	}
+
+	return nil
+}
+
+func (m *Manager) StopStream(streamID uuid.UUID) error {
 	m.DisconnectStream(streamID)
 
 	m.mu.Lock()
@@ -98,7 +158,7 @@ func (m *Manager) StopStream(streamID uuid.UUID) {
 
 	stream, ok := m.clientStreams[streamID]
 	if !ok {
-		return // Stream not found
+		return fmt.Errorf("stream not found: %s", streamID)
 	}
 
 	stream.Timer.Stop()
@@ -121,6 +181,8 @@ func (m *Manager) StopStream(streamID uuid.UUID) {
 			delete(m.providerStreams, id)
 		}
 	}
+
+	return nil
 }
 
 func (m *Manager) ConnectStream(streamID uuid.UUID) (<-chan domain.Message, error) {
@@ -172,7 +234,10 @@ func (m *Manager) DisconnectStream(streamID uuid.UUID) {
 	// Set up the expiry timer
 	stream.Timer = time.AfterFunc(1*time.Minute, func() {
 		fmt.Printf("stream %s expired due to inactivity\n", streamID)
-		m.StopStream(streamID)
+		err := m.StopStream(streamID)
+		if err != nil {
+			fmt.Printf("failed to stop stream after disconnect: %v\n", err)
+		}
 	})
 }
 

services/data_service/internal/provider/binance/futures_websocket.go

@@ -99,7 +99,7 @@ func (b *FuturesWebsocket) IsStreamActive(subject string) bool {
 	return ok
 }
 
-func (b *FuturesWebsocket) Fetch(subject string) (domain.Message, error) {
+func (b *FuturesWebsocket) Fetch(_ string) (domain.Message, error) {
 	return domain.Message{}, fmt.Errorf("not supported: websocket provider does not implement fetch")
 }
 

services/data_service/internal/router/router.go

@@ -2,9 +2,9 @@ package router
 
 import (
 	"fmt"
-	"gitlab.michelsen.id/phillmichelsen/tessera/services/data_service/internal/domain"
 	"sync"
-	"time"
+
+	"gitlab.michelsen.id/phillmichelsen/tessera/services/data_service/internal/domain"
 )
 
 type Router struct {
@@ -15,7 +15,7 @@ type Router struct {
 
 func NewRouter() *Router {
 	return &Router{
-		incoming: make(chan domain.Message, 64), // Buffered channel for incoming messages
+		incoming: make(chan domain.Message, 512), // Buffered channel for incoming messages
 		routes:   make(map[domain.Identifier][]chan<- domain.Message),
 	}
 }
@@ -26,7 +26,6 @@ func (r *Router) IncomingChannel() chan<- domain.Message {
 
 func (r *Router) Run() {
 	for msg := range r.incoming {
-		startTime := time.Now()
 		r.mu.RLock()
 		channels := r.routes[msg.Identifier]
 
@@ -38,7 +37,6 @@ func (r *Router) Run() {
 			}
 		}
 		r.mu.RUnlock()
-		fmt.Printf("Message routed to %d channels in %v\n", len(channels), time.Since(startTime))
 	}
 }
 

services/data_service/internal/server/gprc_control_server.go

@@ -0,0 +1,69 @@
+package server
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/google/uuid"
+	pb "gitlab.michelsen.id/phillmichelsen/tessera/pkg/pb/data_service"
+	"gitlab.michelsen.id/phillmichelsen/tessera/services/data_service/internal/domain"
+	"gitlab.michelsen.id/phillmichelsen/tessera/services/data_service/internal/manager"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+type GRPCControlServer struct {
+	pb.UnimplementedDataServiceControlServer
+	manager *manager.Manager
+}
+
+func NewGRPCControlServer(m *manager.Manager) *GRPCControlServer {
+	return &GRPCControlServer{
+		manager: m,
+	}
+}
+
+func (s *GRPCControlServer) StartStream(_ context.Context, _ *pb.StartStreamRequest) (*pb.StartStreamResponse, error) {
+	streamID, err := s.manager.StartStream()
+	if err != nil {
+		return nil, fmt.Errorf("failed to start stream: %w", err)
+	}
+
+	return &pb.StartStreamResponse{StreamUuid: streamID.String()}, nil
+}
+
+func (s *GRPCControlServer) ConfigureStream(_ context.Context, req *pb.ConfigureStreamRequest) (*pb.ConfigureStreamResponse, error) {
+	streamID, err := uuid.Parse(req.StreamUuid)
+	if err != nil {
+		return nil, status.Errorf(codes.InvalidArgument, "invalid stream_uuid %q: %v", req.StreamUuid, err)
+	}
+
+	// Transform identifiers from protobuf to domain format
+	var ids []domain.Identifier
+	for _, i := range req.Identifiers {
+		ids = append(ids, domain.Identifier{
+			Provider: i.Provider,
+			Subject:  i.Subject,
+		})
+	}
+
+	if err := s.manager.ConfigureStream(streamID, ids); err != nil {
+		return nil, status.Errorf(codes.InvalidArgument, "configure failed: %v", err)
+	}
+
+	return &pb.ConfigureStreamResponse{}, nil
+}
+
+func (s *GRPCControlServer) StopStream(_ context.Context, req *pb.StopStreamRequest) (*pb.StopStreamResponse, error) {
+	streamID, err := uuid.Parse(req.StreamUuid)
+	if err != nil {
+		return nil, status.Errorf(codes.InvalidArgument, "invalid stream_uuid %q: %v", req.StreamUuid, err)
+	}
+
+	err = s.manager.StopStream(streamID) // Should only error if the stream doesn't exist
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "failed to stop stream: %v", err)
+	}
+
+	return &pb.StopStreamResponse{}, nil
+}

services/data_service/internal/server/grpc_streaming_server.go

@@ -1,11 +1,10 @@
 package server
 
 import (
-	"context"
 	"fmt"
+
 	"github.com/google/uuid"
 	pb "gitlab.michelsen.id/phillmichelsen/tessera/pkg/pb/data_service"
-	"gitlab.michelsen.id/phillmichelsen/tessera/services/data_service/internal/domain"
 	"gitlab.michelsen.id/phillmichelsen/tessera/services/data_service/internal/manager"
 )
 
@@ -20,23 +19,6 @@ func NewGRPCStreamingServer(m *manager.Manager) *GRPCStreamingServer {
 	}
 }
 
-func (s *GRPCStreamingServer) StartStream(ctx context.Context, req *pb.StartStreamRequest) (*pb.StartStreamResponse, error) {
-	var ids []domain.Identifier
-	for _, id := range req.Identifiers {
-		ids = append(ids, domain.Identifier{
-			Provider: id.Provider,
-			Subject:  id.Subject,
-		})
-	}
-
-	streamID, err := s.manager.StartStream(ids)
-	if err != nil {
-		return nil, fmt.Errorf("failed to start stream: %w", err)
-	}
-
-	return &pb.StartStreamResponse{StreamUuid: streamID.String()}, nil
-}
-
 func (s *GRPCStreamingServer) ConnectStream(req *pb.ConnectStreamRequest, stream pb.DataServiceStreaming_ConnectStreamServer) error {
 	streamUUID, err := uuid.Parse(req.StreamUuid)
 	if err != nil {
@@ -63,7 +45,7 @@ func (s *GRPCStreamingServer) ConnectStream(req *pb.ConnectStreamRequest, stream
 					Provider: msg.Identifier.Provider,
 					Subject:  msg.Identifier.Subject,
 				},
-				Payload: fmt.Sprintf("%s", msg.Payload),
+				Payload: fmt.Sprintf("%v", msg.Payload),
 			})
 
 			if err != nil {

services/data_service/internal/server/socket_streaming_server.go

@@ -4,11 +4,12 @@ import (
 	"bufio"
 	"encoding/json"
 	"fmt"
-	"github.com/google/uuid"
-	"gitlab.michelsen.id/phillmichelsen/tessera/services/data_service/internal/manager"
 	"io"
 	"net"
 	"strings"
+
+	"github.com/google/uuid"
+	"gitlab.michelsen.id/phillmichelsen/tessera/services/data_service/internal/manager"
 )
 
 type SocketStreamingServer struct {
@@ -34,7 +35,14 @@ func (s *SocketStreamingServer) Serve(lis net.Listener) error {
 }
 
 func (s *SocketStreamingServer) handleConnection(conn net.Conn) {
-	defer conn.Close()
+	defer func(conn net.Conn) {
+		err := conn.Close()
+		if err != nil {
+			fmt.Printf("Failed to close connection: %v\n", err)
+		} else {
+			fmt.Println("Connection closed")
+		}
+	}(conn)
 	reader := bufio.NewReader(conn)
 
 	raw, err := reader.ReadString('\n')