satellite/overlay: do not include gracefully exiting nodes in node selection (#3211)

satellite/overlay/selection_test.go

@@ -6,6 +6,7 @@ package overlay_test
 import (
 	"runtime"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -196,6 +197,110 @@ func testNodeSelection(t *testing.T, ctx *testcontext.Context, planet *testplane
 	}
 }
 
+func TestNodeSelectionGracefulExit(t *testing.T) {
+	testplanet.Run(t, testplanet.Config{
+		SatelliteCount: 1, StorageNodeCount: 10, UplinkCount: 1,
+	}, func(t *testing.T, ctx *testcontext.Context, planet *testplanet.Planet) {
+		satellite := planet.Satellites[0]
+
+		exitingNodes := make(map[storj.NodeID]bool)
+
+		// This sets audit counts of 0, 1, 2, 3, ... 9
+		// so that we can fine-tune how many nodes are considered new or reputable
+		// by modifying the audit count cutoff passed into FindStorageNodesWithPreferences
+		// nodes at indices 0, 2, 4, 6, 8 are gracefully exiting
+		for i, node := range planet.StorageNodes {
+			for k := 0; k < i; k++ {
+				_, err := satellite.DB.OverlayCache().UpdateStats(ctx, &overlay.UpdateRequest{
+					NodeID:       node.ID(),
+					IsUp:         true,
+					AuditSuccess: true,
+					AuditLambda:  1, AuditWeight: 1, AuditDQ: 0.5,
+					UptimeLambda: 1, UptimeWeight: 1, UptimeDQ: 0.5,
+				})
+				require.NoError(t, err)
+			}
+
+			// make half the nodes gracefully exiting
+			if i%2 == 0 {
+				_, err := satellite.DB.OverlayCache().UpdateExitStatus(ctx, &overlay.ExitStatusRequest{
+					NodeID:          node.ID(),
+					ExitInitiatedAt: time.Now(),
+				})
+				require.NoError(t, err)
+				exitingNodes[node.ID()] = true
+			}
+		}
+
+		type test struct {
+			Preferences    overlay.NodeSelectionConfig
+			ExcludeCount   int
+			RequestCount   int
+			ExpectedCount  int
+			ShouldFailWith *errs.Class
+		}
+
+		for i, tt := range []test{
+			{ // reputable and new nodes, happy path
+				Preferences:   testNodeSelectionConfig(5, 0.5, false),
+				RequestCount:  5,
+				ExpectedCount: 5,
+			},
+			{ // all reputable nodes, happy path
+				Preferences:   testNodeSelectionConfig(0, 1, false),
+				RequestCount:  5,
+				ExpectedCount: 5,
+			},
+			{ // all new nodes, happy path
+				Preferences:   testNodeSelectionConfig(50, 1, false),
+				RequestCount:  5,
+				ExpectedCount: 5,
+			},
+			{ // reputable and new nodes, requested too many
+				Preferences:    testNodeSelectionConfig(5, 0.5, false),
+				RequestCount:   10,
+				ExpectedCount:  5,
+				ShouldFailWith: &overlay.ErrNotEnoughNodes,
+			},
+			{ // all reputable nodes, requested too many
+				Preferences:    testNodeSelectionConfig(0, 1, false),
+				RequestCount:   10,
+				ExpectedCount:  5,
+				ShouldFailWith: &overlay.ErrNotEnoughNodes,
+			},
+			{ // all new nodes, requested too many
+				Preferences:    testNodeSelectionConfig(50, 1, false),
+				RequestCount:   10,
+				ExpectedCount:  5,
+				ShouldFailWith: &overlay.ErrNotEnoughNodes,
+			},
+		} {
+			t.Logf("#%2d. %+v", i, tt)
+
+			response, err := satellite.Overlay.Service.FindStorageNodesWithPreferences(ctx, overlay.FindStorageNodesRequest{
+				FreeBandwidth:  0,
+				FreeDisk:       0,
+				RequestedCount: tt.RequestCount,
+			}, &tt.Preferences)
+
+			t.Log(len(response), err)
+			if tt.ShouldFailWith != nil {
+				assert.Error(t, err)
+				assert.True(t, tt.ShouldFailWith.Has(err))
+			} else {
+				assert.NoError(t, err)
+			}
+
+			assert.Equal(t, tt.ExpectedCount, len(response))
+
+			// expect no exiting nodes in selection
+			for _, node := range response {
+				assert.False(t, exitingNodes[node.Id])
+			}
+		}
+	})
+}
+
 func TestDistinctIPs(t *testing.T) {
 	if runtime.GOOS == "darwin" {
 		t.Skip("Test does not work with macOS")

satellite/satellitedb/overlaycache.go

@@ -42,6 +42,7 @@ func (cache *overlaycache) SelectStorageNodes(ctx context.Context, count int, cr
 
 	safeQuery := `
 		WHERE disqualified IS NULL
+		AND exit_initiated_at IS NULL
 		AND type = ?
 		AND free_bandwidth >= ?
 		AND free_disk >= ?
@@ -98,6 +99,7 @@ func (cache *overlaycache) SelectNewStorageNodes(ctx context.Context, count int,
 
 	safeQuery := `
 		WHERE disqualified IS NULL
+		AND exit_initiated_at IS NULL
 		AND type = ?
 		AND free_bandwidth >= ?
 		AND free_disk >= ?