improve jobs endpoint

management/internals/shared/grpc/server.go

@@ -215,12 +215,13 @@ func (s *Server) Job(srv proto.ManagementService_JobServer) error {
 		return status.Errorf(codes.Unauthenticated, "peer is not registered")
 	}
 
-	s.startResponseReceiver(ctx, srv)
-
-	updates := s.jobManager.CreateJobChannel(ctx, accountID, peer.ID)
+	stream := s.jobManager.RegisterStream(ctx, accountID, peer.ID, func(event *job.Event) error {
+		return s.sendJob(ctx, peerKey, event, srv)
+	})
+	defer s.jobManager.UnregisterStream(ctx, accountID, peer.ID, stream)
 	log.WithContext(ctx).Debugf("Job: took %v", time.Since(reqStart))
 
-	return s.sendJobsLoop(ctx, accountID, peerKey, peer, updates, srv)
+	return s.receiveJobResponses(ctx, peerKey, srv)
 }
 
 // Sync validates the existence of a connecting peer, sends an initial state (all available for the connecting peers) and
@@ -362,51 +363,26 @@ func (s *Server) handleHandshake(ctx context.Context, srv proto.ManagementServic
 	return peerKey, nil
 }
 
-func (s *Server) startResponseReceiver(ctx context.Context, srv proto.ManagementService_JobServer) {
-	go func() {
-		for {
-			msg, err := srv.Recv()
-			if err != nil {
-				if errors.Is(err, io.EOF) || errors.Is(err, context.Canceled) {
-					return
-				}
-				log.WithContext(ctx).Warnf("recv job response error: %v", err)
-				return
-			}
-
-			jobResp := &proto.JobResponse{}
-			if _, err := s.parseRequest(ctx, msg, jobResp); err != nil {
-				log.WithContext(ctx).Warnf("invalid job response: %v", err)
-				continue
-			}
-
-			if err := s.jobManager.HandleResponse(ctx, jobResp, msg.WgPubKey); err != nil {
-				log.WithContext(ctx).Errorf("handle job response failed: %v", err)
-			}
-		}
-	}()
-}
-
-func (s *Server) sendJobsLoop(ctx context.Context, accountID string, peerKey wgtypes.Key, peer *nbpeer.Peer, updates *job.Channel, srv proto.ManagementService_JobServer) error {
-	// todo figure out better error handling strategy
-	defer s.jobManager.CloseChannel(ctx, accountID, peer.ID)
-
+func (s *Server) receiveJobResponses(ctx context.Context, peerKey wgtypes.Key, srv proto.ManagementService_JobServer) error {
 	for {
-		event, err := updates.Event(ctx)
+		msg, err := srv.Recv()
 		if err != nil {
-			if errors.Is(err, job.ErrJobChannelClosed) {
-				log.WithContext(ctx).Debugf("jobs channel for peer %s was closed", peerKey.String())
+			if errors.Is(err, io.EOF) || errors.Is(err, context.Canceled) || ctx.Err() != nil {
+				log.WithContext(ctx).Debugf("job stream of peer %s has been closed", peerKey.String())
 				return nil
 			}
-
-			// happens when connection drops, e.g. client disconnects
-			log.WithContext(ctx).Debugf("stream of peer %s has been closed", peerKey.String())
-			return ctx.Err()
+			log.WithContext(ctx).Warnf("recv job response error: %v", err)
+			return err
 		}
 
-		if err := s.sendJob(ctx, peerKey, event, srv); err != nil {
-			log.WithContext(ctx).Warnf("send job failed: %v", err)
-			return nil
+		jobResp := &proto.JobResponse{}
+		if _, err := s.parseRequest(ctx, msg, jobResp); err != nil {
+			log.WithContext(ctx).Warnf("invalid job response: %v", err)
+			continue
+		}
+
+		if err := s.jobManager.HandleResponse(ctx, jobResp, msg.WgPubKey); err != nil {
+			log.WithContext(ctx).Errorf("handle job response failed: %v", err)
 		}
 	}
 }

management/server/job/manager.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"fmt"
 	"sync"
-	"time"
 
 	log "github.com/sirupsen/logrus"
 
@@ -21,11 +20,17 @@ type Event struct {
 	Response *proto.JobResponse
 }
 
+// PeerStream is the send side of a peer's Job stream. Sends are serialized by
+// its mutex; the receive side runs on the stream's gRPC handler goroutine.
+type PeerStream struct {
+	send func(*Event) error
+	mu   sync.Mutex
+}
+
 type Manager struct {
 	mu           *sync.RWMutex
-	jobChannels  map[string]*Channel // per-peer job streams
-	pending      map[string]*Event   // jobID → event
-	responseWait time.Duration
+	streams      map[string]*PeerStream // per-peer job streams
+	pending      map[string]*Event      // jobID → event
 	metrics      telemetry.AppMetrics
 	Store        store.Store
 	peersManager peers.Manager
@@ -34,9 +39,8 @@ type Manager struct {
 func NewJobManager(metrics telemetry.AppMetrics, store store.Store, peersManager peers.Manager) *Manager {
 
 	return &Manager{
-		jobChannels:  make(map[string]*Channel),
+		streams:      make(map[string]*PeerStream),
 		pending:      make(map[string]*Event),
-		responseWait: 5 * time.Minute,
 		metrics:      metrics,
 		mu:           &sync.RWMutex{},
 		Store:        store,
@@ -44,8 +48,9 @@ func NewJobManager(metrics telemetry.AppMetrics, store store.Store, peersManager
 	}
 }
 
-// CreateJobChannel creates or replaces a channel for a peer
-func (jm *Manager) CreateJobChannel(ctx context.Context, accountID, peerID string) *Channel {
+// RegisterStream registers the send side of a peer's Job stream, replacing any
+// previous registration for the peer.
+func (jm *Manager) RegisterStream(ctx context.Context, accountID, peerID string, send func(*Event) error) *PeerStream {
 	// all pending jobs stored in db for this peer should be failed
 	if err := jm.Store.MarkAllPendingJobsAsFailed(ctx, accountID, peerID, "Pending job cleanup: marked as failed automatically due to being stuck too long"); err != nil {
 		log.WithContext(ctx).Error(err.Error())
@@ -54,23 +59,41 @@ func (jm *Manager) CreateJobChannel(ctx context.Context, accountID, peerID strin
 	jm.mu.Lock()
 	defer jm.mu.Unlock()
 
-	if ch, ok := jm.jobChannels[peerID]; ok {
-		ch.Close()
-		delete(jm.jobChannels, peerID)
-	}
+	stream := &PeerStream{send: send}
+	jm.streams[peerID] = stream
+	return stream
+}
 
-	ch := NewChannel()
-	jm.jobChannels[peerID] = ch
-	return ch
+// UnregisterStream removes a peer's stream registration and fails its pending
+// jobs. It is a no-op if the registration was already replaced by a newer
+// stream of the same peer.
+func (jm *Manager) UnregisterStream(ctx context.Context, accountID, peerID string, stream *PeerStream) {
+	jm.mu.Lock()
+	defer jm.mu.Unlock()
+
+	if jm.streams[peerID] != stream {
+		return
+	}
+	delete(jm.streams, peerID)
+
+	for jobID, ev := range jm.pending {
+		if ev.PeerID == peerID {
+			// if the client disconnect and there is pending job then mark it as failed
+			if err := jm.Store.MarkPendingJobsAsFailed(ctx, accountID, peerID, jobID, "Time out peer disconnected"); err != nil {
+				log.WithContext(ctx).Errorf("failed to mark pending jobs as failed: %v", err)
+			}
+			delete(jm.pending, jobID)
+		}
+	}
 }
 
 // SendJob sends a job to a peer and tracks it as pending
 func (jm *Manager) SendJob(ctx context.Context, accountID, peerID string, req *proto.JobRequest) error {
 	jm.mu.RLock()
-	ch, ok := jm.jobChannels[peerID]
+	stream, ok := jm.streams[peerID]
 	jm.mu.RUnlock()
 	if !ok {
-		return fmt.Errorf("peer %s has no channel", peerID)
+		return fmt.Errorf("peer %s has no stream", peerID)
 	}
 
 	event := &Event{
@@ -82,7 +105,10 @@ func (jm *Manager) SendJob(ctx context.Context, accountID, peerID string, req *p
 	jm.pending[string(req.ID)] = event
 	jm.mu.Unlock()
 
-	if err := ch.AddEvent(ctx, jm.responseWait, event); err != nil {
+	stream.mu.Lock()
+	err := stream.send(event)
+	stream.mu.Unlock()
+	if err != nil {
 		jm.cleanup(ctx, accountID, string(req.ID), err.Error())
 		return err
 	}
@@ -127,27 +153,6 @@ func (jm *Manager) HandleResponse(ctx context.Context, resp *proto.JobResponse,
 	return nil
 }
 
-// CloseChannel closes a peer’s channel and cleans up its jobs
-func (jm *Manager) CloseChannel(ctx context.Context, accountID, peerID string) {
-	jm.mu.Lock()
-	defer jm.mu.Unlock()
-
-	if ch, ok := jm.jobChannels[peerID]; ok {
-		ch.Close()
-		delete(jm.jobChannels, peerID)
-	}
-
-	for jobID, ev := range jm.pending {
-		if ev.PeerID == peerID {
-			// if the client disconnect and there is pending job then mark it as failed
-			if err := jm.Store.MarkPendingJobsAsFailed(ctx, accountID, peerID, jobID, "Time out peer disconnected"); err != nil {
-				log.WithContext(ctx).Errorf("failed to mark pending jobs as failed: %v", err)
-			}
-			delete(jm.pending, jobID)
-		}
-	}
-}
-
 // cleanup removes a pending job safely
 func (jm *Manager) cleanup(ctx context.Context, accountID, jobID string, reason string) {
 	jm.mu.Lock()
@@ -165,7 +170,7 @@ func (jm *Manager) IsPeerConnected(peerID string) bool {
 	jm.mu.RLock()
 	defer jm.mu.RUnlock()
 
-	_, ok := jm.jobChannels[peerID]
+	_, ok := jm.streams[peerID]
 	return ok
 }
 

management/server/job/manager_test.go

@@ -0,0 +1,90 @@
+package job
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/management/server/store"
+	"github.com/netbirdio/netbird/shared/management/proto"
+)
+
+func newTestManager(t *testing.T) (*Manager, *store.MockStore) {
+	t.Helper()
+	ctrl := gomock.NewController(t)
+	t.Cleanup(ctrl.Finish)
+	mockStore := store.NewMockStore(ctrl)
+	return NewJobManager(nil, mockStore, nil), mockStore
+}
+
+func TestSendJobDeliversThroughRegisteredStream(t *testing.T) {
+	ctx := context.Background()
+	manager, mockStore := newTestManager(t)
+	mockStore.EXPECT().MarkAllPendingJobsAsFailed(gomock.Any(), "acc", "peer1", gomock.Any()).Return(nil)
+
+	var sent []*Event
+	manager.RegisterStream(ctx, "acc", "peer1", func(event *Event) error {
+		sent = append(sent, event)
+		return nil
+	})
+	require.True(t, manager.IsPeerConnected("peer1"))
+
+	err := manager.SendJob(ctx, "acc", "peer1", &proto.JobRequest{ID: []byte("job1")})
+	require.NoError(t, err)
+	require.Len(t, sent, 1)
+	require.Equal(t, "peer1", sent[0].PeerID)
+	require.True(t, manager.IsPeerHasPendingJobs("peer1"))
+}
+
+func TestSendJobWithoutStream(t *testing.T) {
+	manager, _ := newTestManager(t)
+	err := manager.SendJob(context.Background(), "acc", "peer1", &proto.JobRequest{ID: []byte("job1")})
+	require.Error(t, err)
+}
+
+func TestSendJobFailureCleansPending(t *testing.T) {
+	ctx := context.Background()
+	manager, mockStore := newTestManager(t)
+	mockStore.EXPECT().MarkAllPendingJobsAsFailed(gomock.Any(), "acc", "peer1", gomock.Any()).Return(nil)
+	mockStore.EXPECT().MarkPendingJobsAsFailed(gomock.Any(), "acc", "peer1", "job1", gomock.Any()).Return(nil)
+
+	manager.RegisterStream(ctx, "acc", "peer1", func(*Event) error {
+		return errors.New("stream broken")
+	})
+
+	err := manager.SendJob(ctx, "acc", "peer1", &proto.JobRequest{ID: []byte("job1")})
+	require.Error(t, err)
+	require.False(t, manager.IsPeerHasPendingJobs("peer1"))
+}
+
+func TestUnregisterStreamIgnoresSupersededRegistration(t *testing.T) {
+	ctx := context.Background()
+	manager, mockStore := newTestManager(t)
+	mockStore.EXPECT().MarkAllPendingJobsAsFailed(gomock.Any(), "acc", "peer1", gomock.Any()).Return(nil).Times(2)
+
+	first := manager.RegisterStream(ctx, "acc", "peer1", func(*Event) error { return nil })
+	second := manager.RegisterStream(ctx, "acc", "peer1", func(*Event) error { return nil })
+
+	manager.UnregisterStream(ctx, "acc", "peer1", first)
+	require.True(t, manager.IsPeerConnected("peer1"), "stale unregister must not remove the replacement stream")
+
+	manager.UnregisterStream(ctx, "acc", "peer1", second)
+	require.False(t, manager.IsPeerConnected("peer1"))
+}
+
+func TestUnregisterStreamFailsPendingJobs(t *testing.T) {
+	ctx := context.Background()
+	manager, mockStore := newTestManager(t)
+	mockStore.EXPECT().MarkAllPendingJobsAsFailed(gomock.Any(), "acc", "peer1", gomock.Any()).Return(nil)
+	mockStore.EXPECT().MarkPendingJobsAsFailed(gomock.Any(), "acc", "peer1", "job1", gomock.Any()).Return(nil)
+
+	stream := manager.RegisterStream(ctx, "acc", "peer1", func(*Event) error { return nil })
+	require.NoError(t, manager.SendJob(ctx, "acc", "peer1", &proto.JobRequest{ID: []byte("job1")}))
+	require.True(t, manager.IsPeerHasPendingJobs("peer1"))
+
+	manager.UnregisterStream(ctx, "acc", "peer1", stream)
+	require.False(t, manager.IsPeerHasPendingJobs("peer1"))
+}