feat: add CVM ID to gRPC metadata and implement log-forwarder stream receiving (#611)

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

cmd/agent/main.go

@@ -39,6 +39,7 @@ import (
 	runnerclient "github.com/ultravioletrs/cocos/pkg/clients/grpc/runner"
 	"github.com/ultravioletrs/cocos/pkg/ingress"
 	"golang.org/x/sync/errgroup"
+	"google.golang.org/grpc/metadata"
 )
 
 const (
@@ -170,6 +171,9 @@ func main() {
 		}
 		// Don't defer close here as we want to keep the connection open
 
+		if cfg.CVMId != "" {
+			ctx = metadata.AppendToOutgoingContext(ctx, "job-id", cfg.CVMId)
+		}
 		pc, err := newClient.Process(ctx)
 		if err != nil {
 			grpcClient.Close()
@@ -236,7 +240,11 @@ func main() {
 	}
 	ingressProxy := ingress.NewProxyServer(logger, backendURL, certProvider)
 
-	pc, err := cvmsClient.Process(ctx)
+	agentCtx := ctx
+	if cfg.CVMId != "" {
+		agentCtx = metadata.AppendToOutgoingContext(ctx, "job-id", cfg.CVMId)
+	}
+	pc, err := cvmsClient.Process(agentCtx)
 	if err != nil {
 		logger.Error(fmt.Sprintf("failed to connect to cvm server: %s", err))
 		exitCode = 1

cmd/log-forwarder/main.go

@@ -20,6 +20,7 @@ import (
 	cvmsgrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc/cvm"
 	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
 )
 
 const (
@@ -30,6 +31,7 @@ const (
 
 type config struct {
 	LogLevel string `env:"LOG_FORWARDER_LOG_LEVEL" envAlternate:"AGENT_LOG_LEVEL" envDefault:"debug"`
+	CVMId    string `env:"AGENT_CVM_ID"                 envDefault:""`
 }
 
 func main() {
@@ -100,6 +102,9 @@ func main() {
 	defer cvmClient.Close()
 
 	// Create stream to Manager
+	if cfg.CVMId != "" {
+		ctx = metadata.AppendToOutgoingContext(ctx, "job-id", cfg.CVMId, "connection-type", "log-forwarder")
+	}
 	stream, err := cvmsClient.Process(ctx)
 	if err != nil {
 		logger.Error(fmt.Sprintf("failed to create stream to manager: %s", err))
@@ -122,12 +127,23 @@ func main() {
 			case msg := <-logQueue:
 				if err := stream.Send(msg); err != nil {
 					logger.Error(fmt.Sprintf("failed to send log to manager: %s", err))
-					// Reconnect logic would go here
+					return err
 				}
 			}
 		}
 	})
 
+	// Stream Receiver Goroutine
+	g.Go(func() error {
+		for {
+			_, err := stream.Recv()
+			if err != nil {
+				logger.Error(fmt.Sprintf("stream connection lost: %s", err))
+				return err
+			}
+		}
+	})
+
 	g.Go(func() error {
 		ch := make(chan os.Signal, 1)
 		signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)