From ad314064945f7fcb53cd7d130f154e4db2cb70d5 Mon Sep 17 00:00:00 2001 From: pascal Date: Thu, 16 Jul 2026 02:07:15 +0200 Subject: [PATCH] improve jobs endpoint --- management/internals/shared/grpc/server.go | 62 +++++---------- management/server/job/manager.go | 85 ++++++++++---------- management/server/job/manager_test.go | 90 ++++++++++++++++++++++ 3 files changed, 154 insertions(+), 83 deletions(-) create mode 100644 management/server/job/manager_test.go diff --git a/management/internals/shared/grpc/server.go b/management/internals/shared/grpc/server.go index fa06687d0..cefefac0f 100644 --- a/management/internals/shared/grpc/server.go +++ b/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) } } } diff --git a/management/server/job/manager.go b/management/server/job/manager.go index 0b183ac39..2c3c9220a 100644 --- a/management/server/job/manager.go +++ b/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 } diff --git a/management/server/job/manager_test.go b/management/server/job/manager_test.go new file mode 100644 index 000000000..7a6078c9d --- /dev/null +++ b/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")) +}