[management] Tear down only the peer session that owns the stream (#8057)

management/server/job/manager.go

@@ -57,6 +57,7 @@ func (jm *Manager) CreateJobChannel(ctx context.Context, accountID, peerID strin
 	if ch, ok := jm.jobChannels[peerID]; ok {
 		ch.Close()
 		delete(jm.jobChannels, peerID)
+		jm.failPendingLocked(ctx, accountID, peerID, "Pending job cleanup: job stream replaced by a newer stream")
 	}
 
 	ch := NewChannel()
@@ -127,24 +128,35 @@ 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) {
+// CloseChannel closes the peer's job channel session and fails its pending jobs. It does nothing
+// when a newer job stream has registered a different channel for the peer.
+func (jm *Manager) CloseChannel(ctx context.Context, accountID, peerID string, session *Channel) {
 	jm.mu.Lock()
 	defer jm.mu.Unlock()
 
 	if ch, ok := jm.jobChannels[peerID]; ok {
+		if ch != session {
+			log.WithContext(ctx).Debugf("skipped closing job channel: peer %s is owned by a newer job stream", peerID)
+			return
+		}
 		ch.Close()
 		delete(jm.jobChannels, peerID)
 	}
 
+	jm.failPendingLocked(ctx, accountID, peerID, "Time out peer disconnected")
+}
+
+// failPendingLocked marks the peer's pending jobs as failed and drops them from memory. The caller
+// must hold jm.mu.
+func (jm *Manager) failPendingLocked(ctx context.Context, accountID, peerID, reason string) {
 	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)
+		if ev.PeerID != peerID {
+			continue
 		}
+		if err := jm.Store.MarkPendingJobsAsFailed(ctx, accountID, peerID, jobID, reason); err != nil {
+			log.WithContext(ctx).Errorf("failed to mark pending jobs as failed: %v", err)
+		}
+		delete(jm.pending, jobID)
 	}
 }
 

management/server/job/manager_test.go

@@ -0,0 +1,78 @@
+package job
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/mock/gomock"
+
+	"github.com/netbirdio/netbird/management/server/store"
+)
+
+func TestCloseChannel_OwnSession(t *testing.T) {
+	ctx := context.Background()
+	mockStore := store.NewMockStore(gomock.NewController(t))
+	jm := NewJobManager(nil, mockStore, nil)
+
+	session := NewChannel()
+	jm.jobChannels["peer-1"] = session
+	jm.pending["job-1"] = &Event{PeerID: "peer-1"}
+	jm.pending["job-2"] = &Event{PeerID: "peer-2"}
+
+	mockStore.EXPECT().MarkPendingJobsAsFailed(gomock.Any(), "account-1", "peer-1", "job-1", gomock.Any()).Return(nil)
+
+	jm.CloseChannel(ctx, "account-1", "peer-1", session)
+
+	assert.False(t, jm.IsPeerConnected("peer-1"))
+	_, err := session.Event(ctx)
+	assert.ErrorIs(t, err, ErrJobChannelClosed)
+	assert.NotContains(t, jm.pending, "job-1")
+	assert.Contains(t, jm.pending, "job-2", "jobs of other peers must be kept")
+}
+
+func TestCloseChannel_NewerSessionOwnsPeer(t *testing.T) {
+	ctx := context.Background()
+	jm := NewJobManager(nil, store.NewMockStore(gomock.NewController(t)), nil)
+
+	stale := NewChannel()
+	current := NewChannel()
+	jm.jobChannels["peer-1"] = current
+	jm.pending["job-1"] = &Event{PeerID: "peer-1"}
+
+	jm.CloseChannel(ctx, "account-1", "peer-1", stale)
+
+	require.True(t, jm.IsPeerConnected("peer-1"))
+	assert.Contains(t, jm.pending, "job-1", "jobs served by the newer stream must not be failed")
+
+	event := &Event{PeerID: "peer-1"}
+	require.NoError(t, current.AddEvent(ctx, time.Second, event))
+	got, err := current.Event(ctx)
+	require.NoError(t, err, "newer session channel must stay open")
+	assert.Equal(t, event, got)
+}
+
+func TestCreateJobChannel_ReplacesSessionAndDropsItsPendingJobs(t *testing.T) {
+	ctx := context.Background()
+	mockStore := store.NewMockStore(gomock.NewController(t))
+	jm := NewJobManager(nil, mockStore, nil)
+
+	stale := NewChannel()
+	jm.jobChannels["peer-1"] = stale
+	jm.pending["job-1"] = &Event{PeerID: "peer-1"}
+	jm.pending["job-2"] = &Event{PeerID: "peer-2"}
+
+	mockStore.EXPECT().MarkAllPendingJobsAsFailed(gomock.Any(), "account-1", "peer-1", gomock.Any()).Return(nil)
+	mockStore.EXPECT().MarkPendingJobsAsFailed(gomock.Any(), "account-1", "peer-1", "job-1", "Pending job cleanup: job stream replaced by a newer stream").Return(nil)
+
+	current := jm.CreateJobChannel(ctx, "account-1", "peer-1")
+
+	require.NotSame(t, stale, current)
+	_, err := stale.Event(ctx)
+	assert.ErrorIs(t, err, ErrJobChannelClosed, "replaced channel must be closed")
+	assert.True(t, jm.IsPeerConnected("peer-1"))
+	assert.False(t, jm.IsPeerHasPendingJobs("peer-1"), "jobs of the replaced stream must be dropped")
+	assert.Contains(t, jm.pending, "job-2", "jobs of other peers must be kept")
+}