Move the session cleanup job logic into the session package. (#2060)

* Move the session cleanup job logic into the session package.

This is part of a broader effort to more narrowly define the purpose of the internal/server directory to be for interacting with servers themselves.  Future changes will continue to remove things from this directory until the purpose of the directory is clearly working primarily with the boundary servers resource.

internal/servers/controller/controller.go

@@ -357,28 +357,13 @@ func (c *Controller) registerJobs() error {
 	if err := pluginhost.RegisterJobs(c.baseContext, c.scheduler, rw, rw, c.kms, c.conf.HostPlugins); err != nil {
 		return err
 	}
-
-	if err := c.registerSessionConnectionCleanupJob(); err != nil {
+	if err := session.RegisterJobs(c.baseContext, c.scheduler, rw, c.conf.StatusGracePeriodDuration); err != nil {
 		return err
 	}
 
 	return nil
 }
 
-// registerSessionConnectionCleanupJob is a helper method to abstract
-// registering the session connection cleanup job specifically.
-func (c *Controller) registerSessionConnectionCleanupJob() error {
-	sessionConnectionCleanupJob, err := newSessionConnectionCleanupJob(c.ConnectionRepoFn, c.conf.StatusGracePeriodDuration)
-	if err != nil {
-		return fmt.Errorf("error creating session cleanup job: %w", err)
-	}
-	if err = c.scheduler.RegisterJob(c.baseContext, sessionConnectionCleanupJob); err != nil {
-		return fmt.Errorf("error registering session cleanup job: %w", err)
-	}
-
-	return nil
-}
-
 func (c *Controller) Shutdown() error {
 	const op = "controller.(Controller).Shutdown"
 	if !c.started.Load() {

internal/servers/controller/session_cleanup_job_test.go

@@ -1,170 +0,0 @@
-package controller
-
-import (
-	"context"
-	"fmt"
-	"testing"
-	"time"
-
-	"github.com/hashicorp/boundary/internal/db"
-	"github.com/hashicorp/boundary/internal/errors"
-	"github.com/hashicorp/boundary/internal/iam"
-	"github.com/hashicorp/boundary/internal/kms"
-	"github.com/hashicorp/boundary/internal/scheduler"
-	"github.com/hashicorp/boundary/internal/servers"
-	"github.com/hashicorp/boundary/internal/session"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-)
-
-// assert the interface
-var _ = scheduler.Job(new(sessionConnectionCleanupJob))
-
-// This test has been largely adapted from
-// TestRepository_CloseDeadConnectionsOnWorker in
-// internal/session/repository_connection_test.go.
-func TestSessionConnectionCleanupJob(t *testing.T) {
-	t.Parallel()
-	ctx := context.Background()
-
-	const gracePeriod = 1 * time.Second
-
-	require, assert := require.New(t), assert.New(t)
-	conn, _ := db.TestSetup(t, "postgres")
-	rw := db.New(conn)
-	wrapper := db.TestWrapper(t)
-	iamRepo := iam.TestRepo(t, conn, wrapper)
-	kms := kms.TestKms(t, conn, wrapper)
-	serversRepo, err := servers.NewRepository(rw, rw, kms)
-	require.NoError(err)
-	sessionRepo, err := session.NewRepository(rw, rw, kms)
-	connectionRepo, err := session.NewConnectionRepository(ctx, rw, rw, kms, session.WithDeadWorkerConnCloseMinGrace(gracePeriod))
-	require.NoError(err)
-
-	numConns := 12
-
-	// Create two "workers". One will remain untouched while the other "goes
-	// away and comes back" (worker 2).
-	worker1 := session.TestWorker(t, conn, wrapper, session.WithServerId("worker1"))
-	worker2 := session.TestWorker(t, conn, wrapper, session.WithServerId("worker2"))
-
-	// Create a few sessions on each, activate, and authorize a connection
-	var connIds []string
-	connIdsByWorker := make(map[string][]string)
-	for i := 0; i < numConns; i++ {
-		serverId := worker1.PrivateId
-		if i%2 == 0 {
-			serverId = worker2.PrivateId
-		}
-		sess := session.TestDefaultSession(t, conn, wrapper, iamRepo, session.WithServerId(serverId), session.WithDbOpts(db.WithSkipVetForWrite(true)))
-		sess, _, err = sessionRepo.ActivateSession(ctx, sess.GetPublicId(), sess.Version, serverId, "worker", []byte("foo"))
-		require.NoError(err)
-		c, cs, _, err := session.AuthorizeConnection(ctx, sessionRepo, connectionRepo, sess.GetPublicId(), serverId)
-		require.NoError(err)
-		require.Len(cs, 1)
-		require.Equal(session.StatusAuthorized, cs[0].Status)
-		connIds = append(connIds, c.GetPublicId())
-		if i%2 == 0 {
-			connIdsByWorker[worker2.PrivateId] = append(connIdsByWorker[worker2.PrivateId], c.GetPublicId())
-		} else {
-			connIdsByWorker[worker1.PrivateId] = append(connIdsByWorker[worker1.PrivateId], c.GetPublicId())
-		}
-	}
-
-	// Mark half of the connections connected and leave the others authorized.
-	// This is just to ensure we have a spread when we test it out.
-	for i, connId := range connIds {
-		if i%2 == 0 {
-			_, cs, err := connectionRepo.ConnectConnection(ctx, session.ConnectWith{
-				ConnectionId:       connId,
-				ClientTcpAddress:   "127.0.0.1",
-				ClientTcpPort:      22,
-				EndpointTcpAddress: "127.0.0.1",
-				EndpointTcpPort:    22,
-				UserClientIp:       "127.0.0.1",
-			})
-			require.NoError(err)
-			require.Len(cs, 2)
-			var foundAuthorized, foundConnected bool
-			for _, status := range cs {
-				if status.Status == session.StatusAuthorized {
-					foundAuthorized = true
-				}
-				if status.Status == session.StatusConnected {
-					foundConnected = true
-				}
-			}
-			require.True(foundAuthorized)
-			require.True(foundConnected)
-		}
-	}
-
-	// Create the job.
-	job, err := newSessionConnectionCleanupJob(
-		func() (*session.ConnectionRepository, error) { return connectionRepo, nil },
-		session.DeadWorkerConnCloseMinGrace,
-	)
-	job.gracePeriod = gracePeriod // by-pass factory assert so we dont have to wait so long
-	require.NoError(err)
-
-	// sleep the status grace period.
-	time.Sleep(gracePeriod)
-
-	// Push an upsert to the first worker so that its status has been
-	// updated.
-	_, rowsUpdated, err := serversRepo.UpsertServer(ctx, worker1, []servers.Option{}...)
-	require.NoError(err)
-	require.Equal(1, rowsUpdated)
-
-	// Run the job.
-	require.NoError(job.Run(ctx))
-
-	// Assert connection state on both workers.
-	assertConnections := func(workerId string, closed bool) {
-		connIds, ok := connIdsByWorker[workerId]
-		require.True(ok)
-		require.Len(connIds, 6)
-		for _, connId := range connIds {
-			_, states, err := connectionRepo.LookupConnection(ctx, connId, nil)
-			require.NoError(err)
-			var foundClosed bool
-			for _, state := range states {
-				if state.Status == session.StatusClosed {
-					foundClosed = true
-					break
-				}
-			}
-			assert.Equal(closed, foundClosed)
-		}
-	}
-
-	// Assert that all connections on the second worker are closed
-	assertConnections(worker2.PrivateId, true)
-	// Assert that all connections on the first worker are still open
-	assertConnections(worker1.PrivateId, false)
-}
-
-func TestSessionConnectionCleanupJobNewJobErr(t *testing.T) {
-	t.Parallel()
-	ctx := context.TODO()
-	const op = "controller.newNewSessionConnectionCleanupJob"
-	require := require.New(t)
-
-	job, err := newSessionConnectionCleanupJob(nil, 0)
-	require.Equal(err, errors.E(
-		ctx,
-		errors.WithCode(errors.InvalidParameter),
-		errors.WithOp(op),
-		errors.WithMsg("missing connectionRepoFn"),
-	))
-	require.Nil(job)
-
-	job, err = newSessionConnectionCleanupJob(func() (*session.ConnectionRepository, error) { return nil, nil }, 0)
-	require.Equal(err, errors.E(
-		ctx,
-		errors.WithCode(errors.InvalidParameter),
-		errors.WithOp(op),
-		errors.WithMsg(fmt.Sprintf("invalid gracePeriod, must be greater than %s", session.DeadWorkerConnCloseMinGrace)),
-	))
-	require.Nil(job)
-}

internal/session/jobs.go

@@ -0,0 +1,24 @@
+package session
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/hashicorp/boundary/internal/db"
+	"github.com/hashicorp/boundary/internal/scheduler"
+)
+
+// RegisterJobs registers plugin host related jobs with the provided scheduler.
+func RegisterJobs(ctx context.Context, scheduler *scheduler.Scheduler, w db.Writer, gracePeriod time.Duration) error {
+	const op = "session.RegisterJobs"
+	sessionConnectionCleanupJob, err := newSessionConnectionCleanupJob(w, gracePeriod)
+	if err != nil {
+		return fmt.Errorf("error creating session cleanup job: %w", err)
+	}
+	if err = scheduler.RegisterJob(ctx, sessionConnectionCleanupJob); err != nil {
+		return fmt.Errorf("error registering session cleanup job: %w", err)
+	}
+
+	return nil
+}

internal/session/options.go

@@ -21,24 +21,22 @@ type Option func(*options)
 
 // options = how options are represented
 type options struct {
-	withLimit                       int
-	withOrderByCreateTime           db.OrderBy
-	withScopeIds                    []string
-	withUserId                      string
-	withExpirationTime              *timestamp.Timestamp
-	withTestTofu                    []byte
-	withListingConvert              bool
-	withSessionIds                  []string
-	withServerId                    string
-	withDbOpts                      []db.Option
-	withWorkerStateDelay            time.Duration
-	withDeadWorkerConnCloseMinGrace time.Duration
+	withLimit             int
+	withOrderByCreateTime db.OrderBy
+	withScopeIds          []string
+	withUserId            string
+	withExpirationTime    *timestamp.Timestamp
+	withTestTofu          []byte
+	withListingConvert    bool
+	withSessionIds        []string
+	withServerId          string
+	withDbOpts            []db.Option
+	withWorkerStateDelay  time.Duration
 }
 
 func getDefaultOptions() options {
 	return options{
-		withWorkerStateDelay:            10 * time.Second,
-		withDeadWorkerConnCloseMinGrace: DeadWorkerConnCloseMinGrace,
+		withWorkerStateDelay: 10 * time.Second,
 	}
 }
 
@@ -123,12 +121,3 @@ func WithWorkerStateDelay(d time.Duration) Option {
 		o.withWorkerStateDelay = d
 	}
 }
-
-// WithDeadWorkerConnCloseMinGrace is used to set the minimum allowable setting
-// for the CloseConnectionsForDeadWorkers method. This defaults to the default
-// server liveness setting.
-func WithDeadWorkerConnCloseMinGrace(d time.Duration) Option {
-	return func(o *options) {
-		o.withDeadWorkerConnCloseMinGrace = d
-	}
-}

internal/session/repository_connection.go

@@ -11,16 +11,10 @@ import (
 
 	"github.com/hashicorp/boundary/internal/db"
 	"github.com/hashicorp/boundary/internal/errors"
-	"github.com/hashicorp/boundary/internal/servers"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 )
 
-// DeadWorkerConnCloseMinGrace is the minimum allowable setting for
-// the CloseConnectionsForDeadWorkers method. This is synced with
-// the default server liveness setting.
-var DeadWorkerConnCloseMinGrace = servers.DefaultLiveness
-
 // ConnectionRepository is the session connection database repository.
 type ConnectionRepository struct {
 	reader db.Reader
@@ -33,11 +27,6 @@ type ConnectionRepository struct {
 	// workerStateDelay is used by queries to account for a delay in state propagation between
 	// worker and controller
 	workerStateDelay time.Duration
-
-	// deadWorkerConnCloseMinGrace is the minimum allowable setting for
-	// the CloseConnectionsForDeadWorkers method. This defaults to
-	// the default server liveness setting.
-	deadWorkerConnCloseMinGrace time.Duration
 }
 
 // NewConnectionRepository creates a new session Connection Repository. Supports the options: WithLimit
@@ -60,12 +49,11 @@ func NewConnectionRepository(ctx context.Context, r db.Reader, w db.Writer, kms
 	}
 
 	return &ConnectionRepository{
-		reader:                      r,
-		writer:                      w,
-		kms:                         kms,
-		defaultLimit:                opts.withLimit,
-		workerStateDelay:            opts.withWorkerStateDelay,
-		deadWorkerConnCloseMinGrace: opts.withDeadWorkerConnCloseMinGrace,
+		reader:           r,
+		writer:           w,
+		kms:              kms,
+		defaultLimit:     opts.withLimit,
+		workerStateDelay: opts.withWorkerStateDelay,
 	}, nil
 }
 
@@ -357,59 +345,6 @@ func (r *ConnectionRepository) DeleteConnection(ctx context.Context, publicId st
 	return rowsDeleted, nil
 }
 
-type CloseConnectionsForDeadWorkersResult struct {
-	ServerId                string
-	LastUpdateTime          time.Time
-	NumberConnectionsClosed int
-}
-
-// CloseConnectionsForDeadWorkers will run
-// closeConnectionsForDeadServersCte to look for connections that
-// should be marked because they are on a server that is no longer
-// sending status updates to the controller(s).
-//
-// The only input to the method is the grace period, in seconds.
-func (r *ConnectionRepository) CloseConnectionsForDeadWorkers(ctx context.Context, gracePeriod time.Duration) ([]CloseConnectionsForDeadWorkersResult, error) {
-	const op = "session.(ConnectionRepository).CloseConnectionsForDeadWorkers"
-	if gracePeriod < r.deadWorkerConnCloseMinGrace {
-		return nil, errors.New(ctx,
-			errors.InvalidParameter, op, fmt.Sprintf("gracePeriod must be at least %s", r.deadWorkerConnCloseMinGrace))
-	}
-
-	args := []interface{}{
-		sql.Named("grace_period_seconds", gracePeriod.Seconds()),
-	}
-	results := make([]CloseConnectionsForDeadWorkersResult, 0)
-	_, err := r.writer.DoTx(
-		ctx,
-		db.StdRetryCnt,
-		db.ExpBackoff{},
-		func(reader db.Reader, w db.Writer) error {
-			rows, err := w.Query(ctx, closeConnectionsForDeadServersCte, args)
-			if err != nil {
-				return errors.Wrap(ctx, err, op)
-			}
-			defer rows.Close()
-
-			for rows.Next() {
-				var result CloseConnectionsForDeadWorkersResult
-				if err := w.ScanRows(ctx, rows, &result); err != nil {
-					return errors.Wrap(ctx, err, op)
-				}
-
-				results = append(results, result)
-			}
-
-			return nil
-		},
-	)
-	if err != nil {
-		return nil, errors.Wrap(ctx, err, op)
-	}
-
-	return results, nil
-}
-
 // closeOrphanedConnections looks for connections that are still active, but where not reported by the worker.
 func (r *ConnectionRepository) closeOrphanedConnections(ctx context.Context, workerId string, reportedConnections []string) ([]string, error) {
 	const op = "session.(ConnectionRepository).closeOrphanedConnections"

internal/session/repository_connection_test.go

@@ -11,7 +11,6 @@ import (
 	"github.com/hashicorp/boundary/internal/iam"
 	"github.com/hashicorp/boundary/internal/kms"
 	"github.com/hashicorp/boundary/internal/oplog"
-	"github.com/hashicorp/boundary/internal/servers"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -439,287 +438,6 @@ func TestRepository_orphanedConnections(t *testing.T) {
 	assert.ElementsMatch(found, worker1ConnIds)
 }
 
-func TestRepository_CloseConnectionsForDeadWorkers(t *testing.T) {
-	t.Parallel()
-	ctx := context.Background()
-	require := require.New(t)
-	conn, _ := db.TestSetup(t, "postgres")
-	rw := db.New(conn)
-	wrapper := db.TestWrapper(t)
-	iamRepo := iam.TestRepo(t, conn, wrapper)
-	kms := kms.TestKms(t, conn, wrapper)
-	repo, err := NewRepository(rw, rw, kms)
-	require.NoError(err)
-	deadWorkerConnCloseMinGrace := 1 * time.Second
-	connRepo, err := NewConnectionRepository(ctx, rw, rw, kms, WithDeadWorkerConnCloseMinGrace(deadWorkerConnCloseMinGrace))
-	require.NoError(err)
-	serversRepo, err := servers.NewRepository(rw, rw, kms)
-	require.NoError(err)
-
-	// connection count = 6 * states(authorized, connected, closed = 3) * servers_with_open_connections(3)
-	numConns := 54
-
-	// Create four "workers". This is similar to the setup in
-	// TestRepository_CloseDeadConnectionsOnWorker, but a bit more complex;
-	// firstly, the last worker will have no connections at all, and we will be
-	// closing the others in stages to test multiple servers being closed at
-	// once.
-	worker1 := TestWorker(t, conn, wrapper, WithServerId("worker1"))
-	worker2 := TestWorker(t, conn, wrapper, WithServerId("worker2"))
-	worker3 := TestWorker(t, conn, wrapper, WithServerId("worker3"))
-	worker4 := TestWorker(t, conn, wrapper, WithServerId("worker4"))
-
-	// Create sessions on the first three, activate, and authorize connections
-	var worker1ConnIds, worker2ConnIds, worker3ConnIds []string
-	for i := 0; i < numConns; i++ {
-		var serverId string
-		if i%3 == 0 {
-			serverId = worker1.PrivateId
-		} else if i%3 == 1 {
-			serverId = worker2.PrivateId
-		} else {
-			serverId = worker3.PrivateId
-		}
-		sess := TestDefaultSession(t, conn, wrapper, iamRepo, WithServerId(serverId), WithDbOpts(db.WithSkipVetForWrite(true)))
-		sess, _, err = repo.ActivateSession(ctx, sess.GetPublicId(), sess.Version, serverId, "worker", []byte("foo"))
-		require.NoError(err)
-		c, cs, err := connRepo.AuthorizeConnection(ctx, sess.GetPublicId(), serverId)
-		require.NoError(err)
-		require.Len(cs, 1)
-		require.Equal(StatusAuthorized, cs[0].Status)
-		if i%3 == 0 {
-			worker1ConnIds = append(worker1ConnIds, c.GetPublicId())
-		} else if i%3 == 1 {
-			worker2ConnIds = append(worker2ConnIds, c.GetPublicId())
-		} else {
-			worker3ConnIds = append(worker3ConnIds, c.GetPublicId())
-		}
-	}
-
-	// Mark a third of the connections connected, a third closed, and leave the
-	// others authorized. This is just to ensure we have a spread when we test it
-	// out.
-	for i, connId := range func() []string {
-		var s []string
-		s = append(s, worker1ConnIds...)
-		s = append(s, worker2ConnIds...)
-		s = append(s, worker3ConnIds...)
-		return s
-	}() {
-		if i%3 == 0 {
-			_, cs, err := connRepo.ConnectConnection(ctx, ConnectWith{
-				ConnectionId:       connId,
-				ClientTcpAddress:   "127.0.0.1",
-				ClientTcpPort:      22,
-				EndpointTcpAddress: "127.0.0.1",
-				EndpointTcpPort:    22,
-				UserClientIp:       "127.0.0.1",
-			})
-			require.NoError(err)
-			require.Len(cs, 2)
-			var foundAuthorized, foundConnected bool
-			for _, status := range cs {
-				if status.Status == StatusAuthorized {
-					foundAuthorized = true
-				}
-				if status.Status == StatusConnected {
-					foundConnected = true
-				}
-			}
-			require.True(foundAuthorized)
-			require.True(foundConnected)
-		} else if i%3 == 1 {
-			resp, err := connRepo.closeConnections(ctx, []CloseWith{
-				{
-					ConnectionId: connId,
-					ClosedReason: ConnectionCanceled,
-				},
-			})
-			require.NoError(err)
-			require.Len(resp, 1)
-			cs := resp[0].ConnectionStates
-			require.Len(cs, 2)
-			var foundAuthorized, foundClosed bool
-			for _, status := range cs {
-				if status.Status == StatusAuthorized {
-					foundAuthorized = true
-				}
-				if status.Status == StatusClosed {
-					foundClosed = true
-				}
-			}
-			require.True(foundAuthorized)
-			require.True(foundClosed)
-		}
-	}
-
-	// updateServer is a helper for updating the update time for our
-	// servers. The controller is read back so that we can reference
-	// the most up-to-date fields.
-	updateServer := func(t *testing.T, w *servers.Server) *servers.Server {
-		t.Helper()
-		_, rowsUpdated, err := serversRepo.UpsertServer(ctx, w)
-		require.NoError(err)
-		require.Equal(1, rowsUpdated)
-		servers, err := serversRepo.ListServers(ctx, servers.ServerTypeWorker)
-		require.NoError(err)
-		for _, server := range servers {
-			if server.PrivateId == w.PrivateId {
-				return server
-			}
-		}
-
-		require.FailNowf("server %q not found after updating", w.PrivateId)
-		// Looks weird but needed to build, as we fail in testify instead
-		// of returning an error
-		return nil
-	}
-
-	// requireConnectionStatus is a helper expecting all connections on a worker
-	// to be closed.
-	requireConnectionStatus := func(t *testing.T, connIds []string, expectAllClosed bool) {
-		t.Helper()
-
-		var conns []*Connection
-		require.NoError(repo.list(ctx, &conns, "", nil))
-		for i, connId := range connIds {
-			var expected ConnectionStatus
-			switch {
-			case expectAllClosed:
-				expected = StatusClosed
-
-			case i%3 == 0:
-				expected = StatusConnected
-
-			case i%3 == 1:
-				expected = StatusClosed
-
-			case i%3 == 2:
-				expected = StatusAuthorized
-			}
-
-			_, states, err := connRepo.LookupConnection(ctx, connId)
-			require.NoError(err)
-			require.Equal(expected, states[0].Status, "expected latest status for %q (index %d) to be %v", connId, i, expected)
-		}
-	}
-
-	// We need this helper to fix the zone on protobuf timestamps
-	// versus what gets reported in the
-	// CloseConnectionsForDeadWorkersResult.
-	timestampPbAsUTC := func(t *testing.T, tm time.Time) time.Time {
-		t.Helper()
-		// utcLoc, err := time.LoadLocation("Etc/UTC")
-		// require.NoError(err)
-		return tm.In(time.Local)
-	}
-
-	// Now try some scenarios.
-	{
-		// First, test the error/validation case.
-		result, err := connRepo.CloseConnectionsForDeadWorkers(ctx, -1)
-		require.Equal(err, errors.E(ctx,
-			errors.WithCode(errors.InvalidParameter),
-			errors.WithOp("session.(ConnectionRepository).CloseConnectionsForDeadWorkers"),
-			errors.WithMsg(fmt.Sprintf("gracePeriod must be at least %s", deadWorkerConnCloseMinGrace)),
-		))
-		require.Nil(result)
-	}
-
-	{
-		// Now, try the basis, or where all workers are reporting in.
-		worker1 = updateServer(t, worker1)
-		worker2 = updateServer(t, worker2)
-		worker3 = updateServer(t, worker3)
-		updateServer(t, worker4) // no re-assignment here because we never reference the server again
-
-		result, err := connRepo.CloseConnectionsForDeadWorkers(ctx, deadWorkerConnCloseMinGrace)
-		require.NoError(err)
-		require.Empty(result)
-		// Expect appropriate split connection state on worker1
-		requireConnectionStatus(t, worker1ConnIds, false)
-		// Expect appropriate split connection state on worker2
-		requireConnectionStatus(t, worker2ConnIds, false)
-		// Expect appropriate split connection state on worker3
-		requireConnectionStatus(t, worker3ConnIds, false)
-	}
-
-	{
-		// Now try a zero case - similar to the basis, but only in that no results
-		// are expected to be returned for workers with no connections, even if
-		// they are dead. Here, the server with no connections is worker #4.
-		time.Sleep(deadWorkerConnCloseMinGrace)
-		worker1 = updateServer(t, worker1)
-		worker2 = updateServer(t, worker2)
-		worker3 = updateServer(t, worker3)
-
-		result, err := connRepo.CloseConnectionsForDeadWorkers(ctx, deadWorkerConnCloseMinGrace)
-		require.NoError(err)
-		require.Empty(result)
-		// Expect appropriate split connection state on worker1
-		requireConnectionStatus(t, worker1ConnIds, false)
-		// Expect appropriate split connection state on worker2
-		requireConnectionStatus(t, worker2ConnIds, false)
-		// Expect appropriate split connection state on worker3
-		requireConnectionStatus(t, worker3ConnIds, false)
-	}
-
-	{
-		// The first induction is letting the first worker "die" by not updating it
-		// too. All of its authorized and connected connections should be dead.
-		time.Sleep(deadWorkerConnCloseMinGrace)
-		worker2 = updateServer(t, worker2)
-		worker3 = updateServer(t, worker3)
-
-		result, err := connRepo.CloseConnectionsForDeadWorkers(ctx, deadWorkerConnCloseMinGrace)
-		require.NoError(err)
-		// Assert that we have one result with the appropriate ID and
-		// number of connections closed. Due to how things are
-		require.Equal([]CloseConnectionsForDeadWorkersResult{
-			{
-				ServerId:                worker1.PrivateId,
-				LastUpdateTime:          timestampPbAsUTC(t, worker1.UpdateTime.AsTime()),
-				NumberConnectionsClosed: 12, // 18 per server, with 6 closed already
-			},
-		}, result)
-		// Expect all connections closed on worker1
-		requireConnectionStatus(t, worker1ConnIds, true)
-		// Expect appropriate split connection state on worker2
-		requireConnectionStatus(t, worker2ConnIds, false)
-		// Expect appropriate split connection state on worker3
-		requireConnectionStatus(t, worker3ConnIds, false)
-	}
-
-	{
-		// The final case is having the other two workers die. After
-		// this, we should have all connections closed with the
-		// appropriate message from the next two servers acted on.
-		time.Sleep(deadWorkerConnCloseMinGrace)
-
-		result, err := connRepo.CloseConnectionsForDeadWorkers(ctx, deadWorkerConnCloseMinGrace)
-		require.NoError(err)
-		// Assert that we have one result with the appropriate ID and number of connections closed.
-		require.Equal([]CloseConnectionsForDeadWorkersResult{
-			{
-				ServerId:                worker2.PrivateId,
-				LastUpdateTime:          timestampPbAsUTC(t, worker2.UpdateTime.AsTime()),
-				NumberConnectionsClosed: 12, // 18 per server, with 6 closed already
-			},
-			{
-				ServerId:                worker3.PrivateId,
-				LastUpdateTime:          timestampPbAsUTC(t, worker3.UpdateTime.AsTime()),
-				NumberConnectionsClosed: 12, // 18 per server, with 6 closed already
-			},
-		}, result)
-		// Expect all connections closed on worker1
-		requireConnectionStatus(t, worker1ConnIds, true)
-		// Expect all connections closed on worker2
-		requireConnectionStatus(t, worker2ConnIds, true)
-		// Expect all connections closed on worker3
-		requireConnectionStatus(t, worker3ConnIds, true)
-	}
-}
-
 func TestRepository_CloseConnections(t *testing.T) {
 	t.Parallel()
 	ctx := context.Background()

internal/session/session_cleanup_job.go

@@ -1,18 +1,30 @@
-package controller
+package session
 
 import (
 	"context"
+	"database/sql"
 	stderrors "errors"
 	"fmt"
 	"time"
 
+	"github.com/hashicorp/boundary/internal/db"
 	"github.com/hashicorp/boundary/internal/errors"
 	"github.com/hashicorp/boundary/internal/observability/event"
 	"github.com/hashicorp/boundary/internal/scheduler"
-	"github.com/hashicorp/boundary/internal/servers/controller/common"
-	"github.com/hashicorp/boundary/internal/session"
+	"github.com/hashicorp/boundary/internal/servers"
 )
 
+// deadWorkerConnCloseMinGrace is the minimum allowable setting for
+// the closeConnectionsForDeadWorkers method. This is synced with
+// the default server liveness setting.
+const deadWorkerConnCloseMinGrace = servers.DefaultLiveness
+
+type closeConnectionsForDeadWorkersResult struct {
+	ServerId                string
+	LastUpdateTime          time.Time
+	NumberConnectionsClosed int
+}
+
 // sessionConnectionCleanupJob defines a periodic job that monitors workers for
 // loss of connection and terminates connections on workers that have
 // not sent a heartbeat in a significant period of time.
@@ -23,10 +35,11 @@ import (
 // the controller will win out and order that the connections be
 // closed on the worker.
 type sessionConnectionCleanupJob struct {
-	connectionRepoFn common.ConnectionRepoFactory
+	writer db.Writer
 
 	// The amount of time to give disconnected workers before marking
-	// their connections as closed.
+	// their connections as closed.  This should be at larger than the liveness
+	// setting for the worker.
 	gracePeriod time.Duration
 
 	// The total number of connections closed in the last run.
@@ -35,21 +48,21 @@ type sessionConnectionCleanupJob struct {
 
 // newSessionConnectionCleanupJob instantiates the session cleanup job.
 func newSessionConnectionCleanupJob(
-	connectionRepoFn common.ConnectionRepoFactory,
+	writer db.Writer,
 	gracePeriod time.Duration,
 ) (*sessionConnectionCleanupJob, error) {
-	const op = "controller.newNewSessionConnectionCleanupJob"
+	const op = "session.newNewSessionConnectionCleanupJob"
 	switch {
-	case connectionRepoFn == nil:
-		return nil, errors.NewDeprecated(errors.InvalidParameter, op, "missing connectionRepoFn")
-	case gracePeriod < session.DeadWorkerConnCloseMinGrace:
+	case writer == nil:
+		return nil, errors.NewDeprecated(errors.InvalidParameter, op, "missing db writer")
+	case gracePeriod < deadWorkerConnCloseMinGrace:
 		return nil, errors.NewDeprecated(
-			errors.InvalidParameter, op, fmt.Sprintf("invalid gracePeriod, must be greater than %s", session.DeadWorkerConnCloseMinGrace))
+			errors.InvalidParameter, op, fmt.Sprintf("invalid gracePeriod, must be greater than %s", deadWorkerConnCloseMinGrace))
 	}
 
 	return &sessionConnectionCleanupJob{
-		connectionRepoFn: connectionRepoFn,
-		gracePeriod:      gracePeriod,
+		writer:      writer,
+		gracePeriod: gracePeriod,
 	}, nil
 }
 
@@ -82,17 +95,11 @@ func (j *sessionConnectionCleanupJob) Status() scheduler.JobStatus {
 
 // Run executes the job.
 func (j *sessionConnectionCleanupJob) Run(ctx context.Context) error {
-	const op = "controller.(sessionConnectionCleanupJob).Run"
+	const op = "session.(sessionConnectionCleanupJob).Run"
 	j.totalClosed = 0
 
-	// Load repos.
-	connectionRepo, err := j.connectionRepoFn()
-	if err != nil {
-		return errors.Wrap(ctx, err, op, errors.WithMsg("error getting session repo"))
-	}
-
 	// Run the atomic dead worker cleanup job.
-	results, err := connectionRepo.CloseConnectionsForDeadWorkers(ctx, j.gracePeriod)
+	results, err := j.closeConnectionsForDeadWorkers(ctx, j.gracePeriod)
 	if err != nil {
 		return errors.Wrap(ctx, err, op)
 	}
@@ -111,3 +118,45 @@ func (j *sessionConnectionCleanupJob) Run(ctx context.Context) error {
 
 	return nil
 }
+
+// closeConnectionsForDeadWorkers will run
+// closeConnectionsForDeadServersCte to look for connections that
+// should be marked because they are on a server that is no longer
+// sending status updates to the controller(s).
+//
+// The only input to the method is the grace period, in seconds.
+func (j *sessionConnectionCleanupJob) closeConnectionsForDeadWorkers(ctx context.Context, gracePeriod time.Duration) ([]closeConnectionsForDeadWorkersResult, error) {
+	const op = "session.(sessionConnectionCleanupJob).closeConnectionsForDeadWorkers"
+	args := []interface{}{
+		sql.Named("grace_period_seconds", gracePeriod.Seconds()),
+	}
+	results := make([]closeConnectionsForDeadWorkersResult, 0)
+	_, err := j.writer.DoTx(
+		ctx,
+		db.StdRetryCnt,
+		db.ExpBackoff{},
+		func(reader db.Reader, w db.Writer) error {
+			rows, err := w.Query(ctx, closeConnectionsForDeadServersCte, args)
+			if err != nil {
+				return errors.Wrap(ctx, err, op)
+			}
+			defer rows.Close()
+
+			for rows.Next() {
+				var result closeConnectionsForDeadWorkersResult
+				if err := w.ScanRows(ctx, rows, &result); err != nil {
+					return errors.Wrap(ctx, err, op)
+				}
+
+				results = append(results, result)
+			}
+
+			return nil
+		},
+	)
+	if err != nil {
+		return nil, errors.Wrap(ctx, err, op)
+	}
+
+	return results, nil
+}

internal/session/session_cleanup_job_test.go

@@ -0,0 +1,442 @@
+package session
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/hashicorp/boundary/internal/db"
+	"github.com/hashicorp/boundary/internal/errors"
+	"github.com/hashicorp/boundary/internal/iam"
+	"github.com/hashicorp/boundary/internal/kms"
+	"github.com/hashicorp/boundary/internal/scheduler"
+	"github.com/hashicorp/boundary/internal/servers"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// assert the interface
+var _ = scheduler.Job(new(sessionConnectionCleanupJob))
+
+// This test has been largely adapted from
+// TestRepository_CloseDeadConnectionsOnWorker in
+// internal/session/repository_connection_test.go.
+func TestSessionConnectionCleanupJob(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+
+	const gracePeriod = 1 * time.Second
+
+	require, assert := require.New(t), assert.New(t)
+	conn, _ := db.TestSetup(t, "postgres")
+	rw := db.New(conn)
+	wrapper := db.TestWrapper(t)
+	iamRepo := iam.TestRepo(t, conn, wrapper)
+	kms := kms.TestKms(t, conn, wrapper)
+	serversRepo, err := servers.NewRepository(rw, rw, kms)
+	require.NoError(err)
+	sessionRepo, err := NewRepository(rw, rw, kms)
+	connectionRepo, err := NewConnectionRepository(ctx, rw, rw, kms)
+	require.NoError(err)
+
+	numConns := 12
+
+	// Create two "workers". One will remain untouched while the other "goes
+	// away and comes back" (worker 2).
+	worker1 := TestWorker(t, conn, wrapper, WithServerId("worker1"))
+	worker2 := TestWorker(t, conn, wrapper, WithServerId("worker2"))
+
+	// Create a few sessions on each, activate, and authorize a connection
+	var connIds []string
+	connIdsByWorker := make(map[string][]string)
+	for i := 0; i < numConns; i++ {
+		serverId := worker1.PrivateId
+		if i%2 == 0 {
+			serverId = worker2.PrivateId
+		}
+		sess := TestDefaultSession(t, conn, wrapper, iamRepo, WithServerId(serverId), WithDbOpts(db.WithSkipVetForWrite(true)))
+		sess, _, err = sessionRepo.ActivateSession(ctx, sess.GetPublicId(), sess.Version, serverId, "worker", []byte("foo"))
+		require.NoError(err)
+		c, cs, _, err := AuthorizeConnection(ctx, sessionRepo, connectionRepo, sess.GetPublicId(), serverId)
+		require.NoError(err)
+		require.Len(cs, 1)
+		require.Equal(StatusAuthorized, cs[0].Status)
+		connIds = append(connIds, c.GetPublicId())
+		if i%2 == 0 {
+			connIdsByWorker[worker2.PrivateId] = append(connIdsByWorker[worker2.PrivateId], c.GetPublicId())
+		} else {
+			connIdsByWorker[worker1.PrivateId] = append(connIdsByWorker[worker1.PrivateId], c.GetPublicId())
+		}
+	}
+
+	// Mark half of the connections connected and leave the others authorized.
+	// This is just to ensure we have a spread when we test it out.
+	for i, connId := range connIds {
+		if i%2 == 0 {
+			_, cs, err := connectionRepo.ConnectConnection(ctx, ConnectWith{
+				ConnectionId:       connId,
+				ClientTcpAddress:   "127.0.0.1",
+				ClientTcpPort:      22,
+				EndpointTcpAddress: "127.0.0.1",
+				EndpointTcpPort:    22,
+				UserClientIp:       "127.0.0.1",
+			})
+			require.NoError(err)
+			require.Len(cs, 2)
+			var foundAuthorized, foundConnected bool
+			for _, status := range cs {
+				if status.Status == StatusAuthorized {
+					foundAuthorized = true
+				}
+				if status.Status == StatusConnected {
+					foundConnected = true
+				}
+			}
+			require.True(foundAuthorized)
+			require.True(foundConnected)
+		}
+	}
+
+	// Create the job.
+	job, err := newSessionConnectionCleanupJob(rw, deadWorkerConnCloseMinGrace)
+	job.gracePeriod = gracePeriod // by-pass factory assert so we dont have to wait so long
+	require.NoError(err)
+
+	// sleep the status grace period.
+	time.Sleep(gracePeriod)
+
+	// Push an upsert to the first worker so that its status has been
+	// updated.
+	_, rowsUpdated, err := serversRepo.UpsertServer(ctx, worker1, []servers.Option{}...)
+	require.NoError(err)
+	require.Equal(1, rowsUpdated)
+
+	// Run the job.
+	require.NoError(job.Run(ctx))
+
+	// Assert connection state on both workers.
+	assertConnections := func(workerId string, closed bool) {
+		connIds, ok := connIdsByWorker[workerId]
+		require.True(ok)
+		require.Len(connIds, 6)
+		for _, connId := range connIds {
+			_, states, err := connectionRepo.LookupConnection(ctx, connId, nil)
+			require.NoError(err)
+			var foundClosed bool
+			for _, state := range states {
+				if state.Status == StatusClosed {
+					foundClosed = true
+					break
+				}
+			}
+			assert.Equal(closed, foundClosed)
+		}
+	}
+
+	// Assert that all connections on the second worker are closed
+	assertConnections(worker2.PrivateId, true)
+	// Assert that all connections on the first worker are still open
+	assertConnections(worker1.PrivateId, false)
+}
+
+func TestSessionConnectionCleanupJobNewJobErr(t *testing.T) {
+	t.Parallel()
+	ctx := context.TODO()
+	const op = "session.newNewSessionConnectionCleanupJob"
+	require := require.New(t)
+
+	job, err := newSessionConnectionCleanupJob(nil, 0)
+	require.Equal(err, errors.E(
+		ctx,
+		errors.WithCode(errors.InvalidParameter),
+		errors.WithOp(op),
+		errors.WithMsg("missing db writer"),
+	))
+	require.Nil(job)
+
+	conn, _ := db.TestSetup(t, "postgres")
+	rw := db.New(conn)
+
+	job, err = newSessionConnectionCleanupJob(rw, 0)
+	require.Equal(err, errors.E(
+		ctx,
+		errors.WithCode(errors.InvalidParameter),
+		errors.WithOp(op),
+		errors.WithMsg(fmt.Sprintf("invalid gracePeriod, must be greater than %s", deadWorkerConnCloseMinGrace)),
+	))
+	require.Nil(job)
+}
+
+func TestCloseConnectionsForDeadWorkers(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	require := require.New(t)
+	conn, _ := db.TestSetup(t, "postgres")
+	rw := db.New(conn)
+	wrapper := db.TestWrapper(t)
+	iamRepo := iam.TestRepo(t, conn, wrapper)
+	kms := kms.TestKms(t, conn, wrapper)
+	repo, err := NewRepository(rw, rw, kms)
+	require.NoError(err)
+	gracePeriod := 1 * time.Second
+	connRepo, err := NewConnectionRepository(ctx, rw, rw, kms)
+	require.NoError(err)
+	serversRepo, err := servers.NewRepository(rw, rw, kms)
+	require.NoError(err)
+
+	job, err := newSessionConnectionCleanupJob(rw, deadWorkerConnCloseMinGrace)
+	require.NoError(err)
+
+	// connection count = 6 * states(authorized, connected, closed = 3) * servers_with_open_connections(3)
+	numConns := 54
+
+	// Create four "workers". This is similar to the setup in
+	// TestRepository_CloseDeadConnectionsOnWorker, but a bit more complex;
+	// firstly, the last worker will have no connections at all, and we will be
+	// closing the others in stages to test multiple servers being closed at
+	// once.
+	worker1 := TestWorker(t, conn, wrapper, WithServerId("worker1"))
+	worker2 := TestWorker(t, conn, wrapper, WithServerId("worker2"))
+	worker3 := TestWorker(t, conn, wrapper, WithServerId("worker3"))
+	worker4 := TestWorker(t, conn, wrapper, WithServerId("worker4"))
+
+	// Create sessions on the first three, activate, and authorize connections
+	var worker1ConnIds, worker2ConnIds, worker3ConnIds []string
+	for i := 0; i < numConns; i++ {
+		var serverId string
+		if i%3 == 0 {
+			serverId = worker1.PrivateId
+		} else if i%3 == 1 {
+			serverId = worker2.PrivateId
+		} else {
+			serverId = worker3.PrivateId
+		}
+		sess := TestDefaultSession(t, conn, wrapper, iamRepo, WithServerId(serverId), WithDbOpts(db.WithSkipVetForWrite(true)))
+		sess, _, err = repo.ActivateSession(ctx, sess.GetPublicId(), sess.Version, serverId, "worker", []byte("foo"))
+		require.NoError(err)
+		c, cs, err := connRepo.AuthorizeConnection(ctx, sess.GetPublicId(), serverId)
+		require.NoError(err)
+		require.Len(cs, 1)
+		require.Equal(StatusAuthorized, cs[0].Status)
+		if i%3 == 0 {
+			worker1ConnIds = append(worker1ConnIds, c.GetPublicId())
+		} else if i%3 == 1 {
+			worker2ConnIds = append(worker2ConnIds, c.GetPublicId())
+		} else {
+			worker3ConnIds = append(worker3ConnIds, c.GetPublicId())
+		}
+	}
+
+	// Mark a third of the connections connected, a third closed, and leave the
+	// others authorized. This is just to ensure we have a spread when we test it
+	// out.
+	for i, connId := range func() []string {
+		var s []string
+		s = append(s, worker1ConnIds...)
+		s = append(s, worker2ConnIds...)
+		s = append(s, worker3ConnIds...)
+		return s
+	}() {
+		if i%3 == 0 {
+			_, cs, err := connRepo.ConnectConnection(ctx, ConnectWith{
+				ConnectionId:       connId,
+				ClientTcpAddress:   "127.0.0.1",
+				ClientTcpPort:      22,
+				EndpointTcpAddress: "127.0.0.1",
+				EndpointTcpPort:    22,
+				UserClientIp:       "127.0.0.1",
+			})
+			require.NoError(err)
+			require.Len(cs, 2)
+			var foundAuthorized, foundConnected bool
+			for _, status := range cs {
+				if status.Status == StatusAuthorized {
+					foundAuthorized = true
+				}
+				if status.Status == StatusConnected {
+					foundConnected = true
+				}
+			}
+			require.True(foundAuthorized)
+			require.True(foundConnected)
+		} else if i%3 == 1 {
+			resp, err := connRepo.closeConnections(ctx, []CloseWith{
+				{
+					ConnectionId: connId,
+					ClosedReason: ConnectionCanceled,
+				},
+			})
+			require.NoError(err)
+			require.Len(resp, 1)
+			cs := resp[0].ConnectionStates
+			require.Len(cs, 2)
+			var foundAuthorized, foundClosed bool
+			for _, status := range cs {
+				if status.Status == StatusAuthorized {
+					foundAuthorized = true
+				}
+				if status.Status == StatusClosed {
+					foundClosed = true
+				}
+			}
+			require.True(foundAuthorized)
+			require.True(foundClosed)
+		}
+	}
+
+	// updateServer is a helper for updating the update time for our
+	// servers. The controller is read back so that we can reference
+	// the most up-to-date fields.
+	updateServer := func(t *testing.T, w *servers.Server) *servers.Server {
+		t.Helper()
+		_, rowsUpdated, err := serversRepo.UpsertServer(ctx, w)
+		require.NoError(err)
+		require.Equal(1, rowsUpdated)
+		servers, err := serversRepo.ListServers(ctx, servers.ServerTypeWorker)
+		require.NoError(err)
+		for _, server := range servers {
+			if server.PrivateId == w.PrivateId {
+				return server
+			}
+		}
+
+		require.FailNowf("server %q not found after updating", w.PrivateId)
+		// Looks weird but needed to build, as we fail in testify instead
+		// of returning an error
+		return nil
+	}
+
+	// requireConnectionStatus is a helper expecting all connections on a worker
+	// to be closed.
+	requireConnectionStatus := func(t *testing.T, connIds []string, expectAllClosed bool) {
+		t.Helper()
+
+		var conns []*Connection
+		require.NoError(repo.list(ctx, &conns, "", nil))
+		for i, connId := range connIds {
+			var expected ConnectionStatus
+			switch {
+			case expectAllClosed:
+				expected = StatusClosed
+
+			case i%3 == 0:
+				expected = StatusConnected
+
+			case i%3 == 1:
+				expected = StatusClosed
+
+			case i%3 == 2:
+				expected = StatusAuthorized
+			}
+
+			_, states, err := connRepo.LookupConnection(ctx, connId)
+			require.NoError(err)
+			require.Equal(expected, states[0].Status, "expected latest status for %q (index %d) to be %v", connId, i, expected)
+		}
+	}
+
+	// We need this helper to fix the zone on protobuf timestamps
+	// versus what gets reported in the
+	// closeConnectionsForDeadWorkersResult.
+	timestampPbAsUTC := func(t *testing.T, tm time.Time) time.Time {
+		t.Helper()
+		// utcLoc, err := time.LoadLocation("Etc/UTC")
+		// require.NoError(err)
+		return tm.In(time.Local)
+	}
+
+	// Now try some scenarios.
+	{
+		// Now, try the basis, or where all workers are reporting in.
+		worker1 = updateServer(t, worker1)
+		worker2 = updateServer(t, worker2)
+		worker3 = updateServer(t, worker3)
+		updateServer(t, worker4) // no re-assignment here because we never reference the server again
+
+		result, err := job.closeConnectionsForDeadWorkers(ctx, gracePeriod)
+		require.NoError(err)
+		require.Empty(result)
+		// Expect appropriate split connection state on worker1
+		requireConnectionStatus(t, worker1ConnIds, false)
+		// Expect appropriate split connection state on worker2
+		requireConnectionStatus(t, worker2ConnIds, false)
+		// Expect appropriate split connection state on worker3
+		requireConnectionStatus(t, worker3ConnIds, false)
+	}
+
+	{
+		// Now try a zero case - similar to the basis, but only in that no results
+		// are expected to be returned for workers with no connections, even if
+		// they are dead. Here, the server with no connections is worker #4.
+		time.Sleep(gracePeriod)
+		worker1 = updateServer(t, worker1)
+		worker2 = updateServer(t, worker2)
+		worker3 = updateServer(t, worker3)
+
+		result, err := job.closeConnectionsForDeadWorkers(ctx, gracePeriod)
+		require.NoError(err)
+		require.Empty(result)
+		// Expect appropriate split connection state on worker1
+		requireConnectionStatus(t, worker1ConnIds, false)
+		// Expect appropriate split connection state on worker2
+		requireConnectionStatus(t, worker2ConnIds, false)
+		// Expect appropriate split connection state on worker3
+		requireConnectionStatus(t, worker3ConnIds, false)
+	}
+
+	{
+		// The first induction is letting the first worker "die" by not updating it
+		// too. All of its authorized and connected connections should be dead.
+		time.Sleep(gracePeriod)
+		worker2 = updateServer(t, worker2)
+		worker3 = updateServer(t, worker3)
+
+		result, err := job.closeConnectionsForDeadWorkers(ctx, gracePeriod)
+		require.NoError(err)
+		// Assert that we have one result with the appropriate ID and
+		// number of connections closed. Due to how things are
+		require.Equal([]closeConnectionsForDeadWorkersResult{
+			{
+				ServerId:                worker1.PrivateId,
+				LastUpdateTime:          timestampPbAsUTC(t, worker1.UpdateTime.AsTime()),
+				NumberConnectionsClosed: 12, // 18 per server, with 6 closed already
+			},
+		}, result)
+		// Expect all connections closed on worker1
+		requireConnectionStatus(t, worker1ConnIds, true)
+		// Expect appropriate split connection state on worker2
+		requireConnectionStatus(t, worker2ConnIds, false)
+		// Expect appropriate split connection state on worker3
+		requireConnectionStatus(t, worker3ConnIds, false)
+	}
+
+	{
+		// The final case is having the other two workers die. After
+		// this, we should have all connections closed with the
+		// appropriate message from the next two servers acted on.
+		time.Sleep(gracePeriod)
+
+		result, err := job.closeConnectionsForDeadWorkers(ctx, gracePeriod)
+		require.NoError(err)
+		// Assert that we have one result with the appropriate ID and number of connections closed.
+		require.Equal([]closeConnectionsForDeadWorkersResult{
+			{
+				ServerId:                worker2.PrivateId,
+				LastUpdateTime:          timestampPbAsUTC(t, worker2.UpdateTime.AsTime()),
+				NumberConnectionsClosed: 12, // 18 per server, with 6 closed already
+			},
+			{
+				ServerId:                worker3.PrivateId,
+				LastUpdateTime:          timestampPbAsUTC(t, worker3.UpdateTime.AsTime()),
+				NumberConnectionsClosed: 12, // 18 per server, with 6 closed already
+			},
+		}, result)
+		// Expect all connections closed on worker1
+		requireConnectionStatus(t, worker1ConnIds, true)
+		// Expect all connections closed on worker2
+		requireConnectionStatus(t, worker2ConnIds, true)
+		// Expect all connections closed on worker3
+		requireConnectionStatus(t, worker3ConnIds, true)
+	}
+}