NOISSUE - Fix handling of runreq chunks

.github/workflows/checkproto.yaml

@@ -33,7 +33,7 @@ jobs:
 
       - name: Set up protoc
         run: |
-          PROTOC_VERSION=27.2
+          PROTOC_VERSION=27.3
           PROTOC_GEN_VERSION=v1.34.2
           PROTOC_GRPC_VERSION=v1.4.0
 

manager/api/grpc/client.go

@@ -3,9 +3,9 @@
 package grpc
 
 import (
-	"bytes"
 	"context"
 	"log/slog"
+	"sync"
 	"time"
 
 	"github.com/absmach/magistrala/pkg/errors"
@@ -22,19 +22,21 @@ var (
 )
 
 type ManagerClient struct {
-	stream       pkgmanager.ManagerService_ProcessClient
-	svc          manager.Service
-	messageQueue chan *pkgmanager.ClientStreamMessage
-	logger       *slog.Logger
+	stream        pkgmanager.ManagerService_ProcessClient
+	svc           manager.Service
+	messageQueue  chan *pkgmanager.ClientStreamMessage
+	logger        *slog.Logger
+	runReqManager *runRequestManager
 }
 
 // NewClient returns new gRPC client instance.
 func NewClient(stream pkgmanager.ManagerService_ProcessClient, svc manager.Service, messageQueue chan *pkgmanager.ClientStreamMessage, logger *slog.Logger) ManagerClient {
 	return ManagerClient{
-		stream:       stream,
-		svc:          svc,
-		messageQueue: messageQueue,
-		logger:       logger,
+		stream:        stream,
+		svc:           svc,
+		messageQueue:  messageQueue,
+		logger:        logger,
+		runReqManager: newRunRequestManager(),
 	}
 }
 
@@ -53,7 +55,6 @@ func (client ManagerClient) Process(ctx context.Context, cancel context.CancelFu
 }
 
 func (client ManagerClient) handleIncomingMessages(ctx context.Context) error {
-	var runReqBuffer bytes.Buffer
 	for {
 		select {
 		case <-ctx.Done():
@@ -63,39 +64,42 @@ func (client ManagerClient) handleIncomingMessages(ctx context.Context) error {
 			if err != nil {
 				return err
 			}
-			if err := client.processIncomingMessage(ctx, req, &runReqBuffer); err != nil {
+			if err := client.processIncomingMessage(ctx, req); err != nil {
 				return err
 			}
 		}
 	}
 }
 
-func (client ManagerClient) processIncomingMessage(ctx context.Context, req *pkgmanager.ServerStreamMessage, runReqBuffer *bytes.Buffer) error {
+func (client ManagerClient) processIncomingMessage(ctx context.Context, req *pkgmanager.ServerStreamMessage) error {
 	switch mes := req.Message.(type) {
 	case *pkgmanager.ServerStreamMessage_RunReqChunks:
-		return client.handleRunReqChunks(ctx, mes, runReqBuffer)
+		return client.handleRunReqChunks(ctx, mes)
 	case *pkgmanager.ServerStreamMessage_TerminateReq:
 		return client.handleTerminateReq(mes)
 	case *pkgmanager.ServerStreamMessage_StopComputation:
 		go client.handleStopComputation(ctx, mes)
 	case *pkgmanager.ServerStreamMessage_BackendInfoReq:
-		go client.handleBackendInfoReq(ctx, mes)
+		go client.handleBackendInfoReq(mes)
 	default:
 		return errors.New("unknown message type")
 	}
 	return nil
 }
 
-func (client ManagerClient) handleRunReqChunks(ctx context.Context, mes *pkgmanager.ServerStreamMessage_RunReqChunks, runReqBuffer *bytes.Buffer) error {
-	if len(mes.RunReqChunks.Data) == 0 {
+func (client *ManagerClient) handleRunReqChunks(ctx context.Context, mes *pkgmanager.ServerStreamMessage_RunReqChunks) error {
+	buffer, complete := client.runReqManager.addChunk(mes.RunReqChunks.Id, mes.RunReqChunks.Data, mes.RunReqChunks.IsLast)
+
+	if complete {
 		var runReq pkgmanager.ComputationRunReq
-		if err := proto.Unmarshal(runReqBuffer.Bytes(), &runReq); err != nil {
+		if err := proto.Unmarshal(buffer, &runReq); err != nil {
 			return errors.Wrap(err, errCorruptedManifest)
 		}
+
 		go client.executeRun(ctx, &runReq)
 	}
-	_, err := runReqBuffer.Write(mes.RunReqChunks.Data)
-	return err
+
+	return nil
 }
 
 func (client ManagerClient) executeRun(ctx context.Context, runReq *pkgmanager.ComputationRunReq) {
@@ -129,7 +133,7 @@ func (client ManagerClient) handleStopComputation(ctx context.Context, mes *pkgm
 	client.sendMessage(&pkgmanager.ClientStreamMessage{Message: msg})
 }
 
-func (client ManagerClient) handleBackendInfoReq(ctx context.Context, mes *pkgmanager.ServerStreamMessage_BackendInfoReq) {
+func (client ManagerClient) handleBackendInfoReq(mes *pkgmanager.ServerStreamMessage_BackendInfoReq) {
 	res, err := client.svc.FetchBackendInfo()
 	if err != nil {
 		client.logger.Warn(err.Error())
@@ -167,3 +171,55 @@ func (client ManagerClient) sendMessage(mes *pkgmanager.ClientStreamMessage) {
 		client.logger.Warn("Failed to send message: timeout exceeded")
 	}
 }
+
+type runRequestManager struct {
+	requests map[string]*runRequest
+	mu       sync.Mutex
+}
+
+type runRequest struct {
+	buffer    []byte
+	lastChunk time.Time
+	timer     *time.Timer
+}
+
+func newRunRequestManager() *runRequestManager {
+	return &runRequestManager{
+		requests: make(map[string]*runRequest),
+	}
+}
+
+func (m *runRequestManager) addChunk(id string, chunk []byte, isLast bool) ([]byte, bool) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	req, exists := m.requests[id]
+	if !exists {
+		req = &runRequest{
+			buffer:    make([]byte, 0),
+			lastChunk: time.Now(),
+			timer:     time.AfterFunc(runReqTimeout, func() { m.timeoutRequest(id) }),
+		}
+		m.requests[id] = req
+	}
+
+	req.buffer = append(req.buffer, chunk...)
+	req.lastChunk = time.Now()
+	req.timer.Reset(runReqTimeout)
+
+	if isLast {
+		delete(m.requests, id)
+		req.timer.Stop()
+		return req.buffer, true
+	}
+
+	return nil, false
+}
+
+func (m *runRequestManager) timeoutRequest(id string) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	delete(m.requests, id)
+	// Log timeout or handle it as needed
+}

manager/api/grpc/server.go

@@ -7,6 +7,7 @@ import (
 	"context"
 	"errors"
 	"io"
+	"time"
 
 	"github.com/ultravioletrs/cocos/pkg/manager"
 	"golang.org/x/sync/errgroup"
@@ -20,7 +21,10 @@ var (
 	ErrUnexpectedMsg                              = errors.New("unknown message type")
 )
 
-const bufferSize = 1024 * 1024 // 1 MB
+const (
+	bufferSize    = 1024 * 1024 // 1 MB
+	runReqTimeout = 30 * time.Second
+)
 
 type SendFunc func(*manager.ServerStreamMessage) error
 
@@ -89,14 +93,20 @@ func (s *grpcServer) sendRunReqInChunks(stream manager.ManagerService_ProcessSer
 
 	for {
 		n, err := dataBuffer.Read(buf)
-		if err != nil && err != io.EOF {
+		isLast := false
+
+		if err == io.EOF {
+			isLast = true
+		} else if err != nil {
 			return err
 		}
 
 		chunk := &manager.ServerStreamMessage{
 			Message: &manager.ServerStreamMessage_RunReqChunks{
 				RunReqChunks: &manager.RunReqChunks{
-					Data: buf[:n],
+					Id:     runReq.Id,
+					Data:   buf[:n],
+					IsLast: isLast,
 				},
 			},
 		}
@@ -105,7 +115,7 @@ func (s *grpcServer) sendRunReqInChunks(stream manager.ManagerService_ProcessSer
 			return err
 		}
 
-		if err == io.EOF {
+		if isLast {
 			break
 		}
 	}

manager/manager.proto

@@ -74,6 +74,8 @@ message ServerStreamMessage {
 
 message RunReqChunks {
   bytes data = 1;
+  string id = 2;
+  bool is_last = 3;
 }
 
 message ComputationRunReq {

manager/qemu/vm.go

@@ -8,6 +8,7 @@ import (
 	"os/exec"
 
 	"github.com/gofrs/uuid"
+	"github.com/ultravioletrs/cocos/internal"
 	"github.com/ultravioletrs/cocos/manager/vm"
 	"github.com/ultravioletrs/cocos/pkg/manager"
 )
@@ -16,6 +17,7 @@ const (
 	firmwareVars = "OVMF_VARS"
 	KernelFile   = "bzImage"
 	rootfsFile   = "rootfs.cpio"
+	tmpDir       = "/tmp"
 )
 
 type qemuVM struct {
@@ -43,6 +45,17 @@ func (v *qemuVM) Start() error {
 	v.config.NetDevConfig.ID = fmt.Sprintf("%s-%s", v.config.NetDevConfig.ID, id)
 	v.config.SevConfig.ID = fmt.Sprintf("%s-%s", v.config.SevConfig.ID, id)
 
+	if !v.config.KernelHash {
+		// Copy firmware vars file.
+		srcFile := v.config.OVMFVarsConfig.File
+		dstFile := fmt.Sprintf("%s/%s-%s.fd", tmpDir, firmwareVars, id)
+		err = internal.CopyFile(srcFile, dstFile)
+		if err != nil {
+			return err
+		}
+		v.config.OVMFVarsConfig.File = dstFile
+	}
+
 	exe, args, err := v.executableAndArgs()
 	if err != nil {
 		return err