Graceful shutdown blocks on running proxy handler connections (#2789)

internal/server/options.go

@@ -48,6 +48,7 @@ type options struct {
 	withActiveWorkers                      bool
 	withFeature                            version.Feature
 	withDirectlyConnected                  bool
+	withWorkerPool                         []string
 }
 
 func getDefaultOptions() options {
@@ -238,3 +239,10 @@ func WithDirectlyConnected(conn bool) Option {
 		o.withDirectlyConnected = conn
 	}
 }
+
+// WithWorkerPool provides a slice of worker ids.
+func WithWorkerPool(workerIds []string) Option {
+	return func(o *options) {
+		o.withWorkerPool = workerIds
+	}
+}

internal/server/options_test.go

@@ -225,4 +225,12 @@ func Test_GetOpts(t *testing.T) {
 		testOpts.withNewIdFunc = nil
 		assert.Equal(t, opts, testOpts)
 	})
+	t.Run("WithWorkerPool", func(t *testing.T) {
+		opts := GetOpts(WithWorkerPool([]string{"1", "2", "3"}))
+		testOpts := getDefaultOptions()
+		testOpts.withWorkerPool = []string{"1", "2", "3"}
+		opts.withNewIdFunc = nil
+		testOpts.withNewIdFunc = nil
+		assert.Equal(t, opts, testOpts)
+	})
 }