Chris Marchesi
5a70875726
This commit adds the support to do the following: * Mark connections for non-reporting workers as closed. This is the controller counterpart to the worker functionality (see #1330). This is written as a scheduled job that does the work DB-side in a single atomic query. * Works to reconcile states if such a broken controller-worker connection resumes and a worker reports a connection as connected that should be disconnected. In this case, the controller will send an update request, and the worker will honor it and terminate the connection. * Further refinement of the grace period setting has been added here. We have converged on the current server "liveness" setting as our default here, which is half of the previous 30s (15 seconds, in other words). Additionally, this is now configurable on the controller and worker side, with the caveat that it's currently impossible to do so in config as the setting has been untagged in HCL. This is exposed so that we can run some sophisticated testing scenarios where we skew the grace period to either the controller or worker to ensure the aforementioned reconciliation works. * Some repository functions have been added to support the new functionality, in addition to some test code to the worker to allow querying of session state while testing. * Finally, we've added some add timestamp subtraction functions as well, basically serving as the opposite of the addition functions.
772 lines
24 KiB
Go
772 lines
24 KiB
Go
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/oplog"
|
|
"github.com/hashicorp/boundary/internal/servers"
|
|
"github.com/hashicorp/shared-secure-libs/strutil"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestRepository_ListConnection(t *testing.T) {
|
|
t.Parallel()
|
|
conn, _ := db.TestSetup(t, "postgres")
|
|
const testLimit = 10
|
|
wrapper := db.TestWrapper(t)
|
|
iamRepo := iam.TestRepo(t, conn, wrapper)
|
|
rw := db.New(conn)
|
|
kms := kms.TestKms(t, conn, wrapper)
|
|
repo, err := NewRepository(rw, rw, kms, WithLimit(testLimit))
|
|
require.NoError(t, err)
|
|
session := TestDefaultSession(t, conn, wrapper, iamRepo)
|
|
|
|
type args struct {
|
|
searchForSessionId string
|
|
opt []Option
|
|
}
|
|
tests := []struct {
|
|
name string
|
|
createCnt int
|
|
args args
|
|
wantCnt int
|
|
wantErr bool
|
|
}{
|
|
{
|
|
name: "no-limit",
|
|
createCnt: repo.defaultLimit + 1,
|
|
args: args{
|
|
searchForSessionId: session.PublicId,
|
|
opt: []Option{WithLimit(-1)},
|
|
},
|
|
wantCnt: repo.defaultLimit + 1,
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "default-limit",
|
|
createCnt: repo.defaultLimit + 1,
|
|
args: args{
|
|
searchForSessionId: session.PublicId,
|
|
},
|
|
wantCnt: repo.defaultLimit,
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "custom-limit",
|
|
createCnt: repo.defaultLimit + 1,
|
|
args: args{
|
|
searchForSessionId: session.PublicId,
|
|
opt: []Option{WithLimit(3)},
|
|
},
|
|
wantCnt: 3,
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "bad-session-id",
|
|
createCnt: repo.defaultLimit + 1,
|
|
args: args{
|
|
searchForSessionId: "s_thisIsNotValid",
|
|
},
|
|
wantCnt: 0,
|
|
wantErr: false,
|
|
},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
assert, require := assert.New(t), require.New(t)
|
|
require.NoError(conn.Where("1=1").Delete(AllocConnection()).Error)
|
|
testConnections := []*Connection{}
|
|
for i := 0; i < tt.createCnt; i++ {
|
|
c := TestConnection(t, conn,
|
|
session.PublicId,
|
|
"127.0.0.1",
|
|
22,
|
|
"127.0.0.1",
|
|
2222,
|
|
)
|
|
testConnections = append(testConnections, c)
|
|
}
|
|
assert.Equal(tt.createCnt, len(testConnections))
|
|
got, err := repo.ListConnectionsBySessionId(context.Background(), tt.args.searchForSessionId, tt.args.opt...)
|
|
if tt.wantErr {
|
|
require.Error(err)
|
|
return
|
|
}
|
|
require.NoError(err)
|
|
assert.Equal(tt.wantCnt, len(got))
|
|
})
|
|
}
|
|
t.Run("withOrder", func(t *testing.T) {
|
|
assert, require := assert.New(t), require.New(t)
|
|
require.NoError(conn.Where("1=1").Delete(AllocConnection()).Error)
|
|
wantCnt := 5
|
|
for i := 0; i < wantCnt; i++ {
|
|
_ = TestConnection(t, conn,
|
|
session.PublicId,
|
|
"127.0.0.1",
|
|
22,
|
|
"127.0.0.1",
|
|
2222,
|
|
)
|
|
}
|
|
got, err := repo.ListConnectionsBySessionId(context.Background(), session.PublicId, WithOrderByCreateTime(db.AscendingOrderBy))
|
|
require.NoError(err)
|
|
assert.Equal(wantCnt, len(got))
|
|
|
|
for i := 0; i < len(got)-1; i++ {
|
|
first := got[i].CreateTime.Timestamp.AsTime()
|
|
second := got[i+1].CreateTime.Timestamp.AsTime()
|
|
assert.True(first.Before(second))
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestRepository_DeleteConnection(t *testing.T) {
|
|
t.Parallel()
|
|
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(t, err)
|
|
session := TestDefaultSession(t, conn, wrapper, iamRepo)
|
|
|
|
type args struct {
|
|
connection *Connection
|
|
opt []Option
|
|
}
|
|
tests := []struct {
|
|
name string
|
|
args args
|
|
wantRowsDeleted int
|
|
wantErr bool
|
|
wantErrMsg string
|
|
}{
|
|
{
|
|
name: "valid",
|
|
args: args{
|
|
connection: TestConnection(t, conn, session.PublicId, "127.0.0.1", 22, "127.0.0.1", 2222),
|
|
},
|
|
wantRowsDeleted: 1,
|
|
wantErr: false,
|
|
},
|
|
{
|
|
name: "no-public-id",
|
|
args: args{
|
|
connection: func() *Connection {
|
|
c := AllocConnection()
|
|
return &c
|
|
}(),
|
|
},
|
|
wantRowsDeleted: 0,
|
|
wantErr: true,
|
|
wantErrMsg: "session.(Repository).DeleteConnection: missing public id: parameter violation: error #100",
|
|
},
|
|
{
|
|
name: "not-found",
|
|
args: args{
|
|
connection: func() *Connection {
|
|
c := AllocConnection()
|
|
id, err := newConnectionId()
|
|
require.NoError(t, err)
|
|
c.PublicId = id
|
|
return &c
|
|
}(),
|
|
},
|
|
wantRowsDeleted: 0,
|
|
wantErr: true,
|
|
wantErrMsg: "db.LookupById: record not found",
|
|
},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
assert := assert.New(t)
|
|
deletedRows, err := repo.DeleteConnection(context.Background(), tt.args.connection.PublicId, tt.args.opt...)
|
|
if tt.wantErr {
|
|
assert.Error(err)
|
|
assert.Equal(0, deletedRows)
|
|
assert.Contains(err.Error(), tt.wantErrMsg)
|
|
err = db.TestVerifyOplog(t, rw, tt.args.connection.PublicId, db.WithOperation(oplog.OpType_OP_TYPE_DELETE), db.WithCreateNotBefore(10*time.Second))
|
|
assert.Error(err)
|
|
assert.True(errors.IsNotFoundError(err))
|
|
return
|
|
}
|
|
assert.NoError(err)
|
|
assert.Equal(tt.wantRowsDeleted, deletedRows)
|
|
found, _, err := repo.LookupConnection(context.Background(), tt.args.connection.PublicId)
|
|
assert.NoError(err)
|
|
assert.Nil(found)
|
|
|
|
err = db.TestVerifyOplog(t, rw, tt.args.connection.PublicId, db.WithOperation(oplog.OpType_OP_TYPE_DELETE), db.WithCreateNotBefore(10*time.Second))
|
|
assert.Error(err)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestRepository_CloseDeadConnectionsOnWorker(t *testing.T) {
|
|
t.Parallel()
|
|
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)
|
|
repo, err := NewRepository(rw, rw, kms)
|
|
require.NoError(err)
|
|
ctx := context.Background()
|
|
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
|
|
var worker2ConnIds []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 = repo.ActivateSession(ctx, sess.GetPublicId(), sess.Version, serverId, "worker", []byte("foo"))
|
|
require.NoError(err)
|
|
c, cs, _, err := repo.AuthorizeConnection(ctx, 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 {
|
|
worker2ConnIds = append(worker2ConnIds, 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 := repo.ConnectConnection(ctx, ConnectWith{
|
|
ConnectionId: connId,
|
|
ClientTcpAddress: "127.0.0.1",
|
|
ClientTcpPort: 22,
|
|
EndpointTcpAddress: "127.0.0.1",
|
|
EndpointTcpPort: 22,
|
|
})
|
|
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)
|
|
}
|
|
}
|
|
|
|
// There is a 10 second delay to account for time for the connections to
|
|
// transition
|
|
time.Sleep(15 * time.Second)
|
|
|
|
// Now, advertise only some of the connection IDs for worker 2. After,
|
|
// all connection IDs for worker 1 should be showing as non-closed, and
|
|
// the ones for worker 2 not advertised should be closed.
|
|
shouldStayOpen := worker2ConnIds[0:2]
|
|
count, err := repo.CloseDeadConnectionsForWorker(ctx, worker2.GetPrivateId(), shouldStayOpen)
|
|
require.NoError(err)
|
|
assert.Equal(4, count)
|
|
|
|
// For the ones we didn't specify, we expect those to now be closed. We
|
|
// expect all others to be open.
|
|
|
|
shouldBeClosed := worker2ConnIds[2:]
|
|
var conns []*Connection
|
|
require.NoError(repo.list(ctx, &conns, "", nil))
|
|
for _, conn := range conns {
|
|
_, states, err := repo.LookupConnection(ctx, conn.PublicId)
|
|
require.NoError(err)
|
|
var foundClosed bool
|
|
for _, state := range states {
|
|
if state.Status == StatusClosed {
|
|
foundClosed = true
|
|
break
|
|
}
|
|
}
|
|
assert.True(foundClosed == strutil.StrListContains(shouldBeClosed, conn.PublicId))
|
|
}
|
|
|
|
// Now, advertise none of the connection IDs for worker 2. This is mainly to
|
|
// test that handling the case where we do not include IDs works properly as
|
|
// it changes the where clause.
|
|
count, err = repo.CloseDeadConnectionsForWorker(ctx, worker1.GetPrivateId(), nil)
|
|
require.NoError(err)
|
|
assert.Equal(6, count)
|
|
|
|
// We now expect all but those blessed few to be closed
|
|
conns = nil
|
|
require.NoError(repo.list(ctx, &conns, "", nil))
|
|
for _, conn := range conns {
|
|
_, states, err := repo.LookupConnection(ctx, conn.PublicId)
|
|
require.NoError(err)
|
|
var foundClosed bool
|
|
for _, state := range states {
|
|
if state.Status == StatusClosed {
|
|
foundClosed = true
|
|
break
|
|
}
|
|
}
|
|
assert.True(foundClosed != strutil.StrListContains(shouldStayOpen, conn.PublicId))
|
|
}
|
|
}
|
|
|
|
func TestRepository_CloseConnectionsForDeadWorkers(t *testing.T) {
|
|
t.Parallel()
|
|
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)
|
|
serversRepo, err := servers.NewRepository(rw, rw, kms)
|
|
require.NoError(err)
|
|
ctx := context.Background()
|
|
|
|
// 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 := repo.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 := repo.ConnectConnection(ctx, ConnectWith{
|
|
ConnectionId: connId,
|
|
ClientTcpAddress: "127.0.0.1",
|
|
ClientTcpPort: 22,
|
|
EndpointTcpAddress: "127.0.0.1",
|
|
EndpointTcpPort: 22,
|
|
})
|
|
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 := repo.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)
|
|
}
|
|
}
|
|
|
|
// There is a 15 second delay to account for time for the connections to
|
|
// transition
|
|
time.Sleep(15 * time.Second)
|
|
|
|
// 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 := repo.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(utcLoc)
|
|
}
|
|
|
|
// Now try some scenarios.
|
|
{
|
|
// First, test the error/validation case.
|
|
result, err := repo.CloseConnectionsForDeadWorkers(ctx, 0)
|
|
require.Equal(err, errors.E(
|
|
errors.WithCode(errors.InvalidParameter),
|
|
errors.WithOp("session.(Repository).CloseConnectionsForDeadWorkers"),
|
|
errors.WithMsg(fmt.Sprintf("gracePeriod must be at least %d seconds", 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 := repo.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(time.Second * time.Duration(DeadWorkerConnCloseMinGrace))
|
|
worker1 = updateServer(t, worker1)
|
|
worker2 = updateServer(t, worker2)
|
|
worker3 = updateServer(t, worker3)
|
|
|
|
result, err := repo.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(time.Second * time.Duration(DeadWorkerConnCloseMinGrace))
|
|
worker2 = updateServer(t, worker2)
|
|
worker3 = updateServer(t, worker3)
|
|
|
|
result, err := repo.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(time.Second * time.Duration(DeadWorkerConnCloseMinGrace))
|
|
|
|
result, err := repo.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_ShouldCloseConnectionsOnWorker(t *testing.T) {
|
|
t.Parallel()
|
|
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)
|
|
ctx := context.Background()
|
|
numConns := 12
|
|
|
|
// Create a worker, we only need one here as our query is dependent
|
|
// on connection and not worker.
|
|
worker1 := TestWorker(t, conn, wrapper, WithServerId("worker1"))
|
|
|
|
// Create a few sessions on each, activate, and authorize a connection
|
|
var connIds []string
|
|
sessionConnIds := make(map[string][]string)
|
|
for i := 0; i < numConns; i++ {
|
|
serverId := worker1.PrivateId
|
|
sess := TestDefaultSession(t, conn, wrapper, iamRepo, WithServerId(serverId), WithDbOpts(db.WithSkipVetForWrite(true)))
|
|
sessionId := sess.GetPublicId()
|
|
sess, _, err = repo.ActivateSession(ctx, sessionId, sess.Version, serverId, "worker", []byte("foo"))
|
|
require.NoError(err)
|
|
c, cs, _, err := repo.AuthorizeConnection(ctx, sess.GetPublicId(), serverId)
|
|
require.NoError(err)
|
|
require.Len(cs, 1)
|
|
require.Equal(StatusAuthorized, cs[0].Status)
|
|
connId := c.GetPublicId()
|
|
connIds = append(connIds, connId)
|
|
sessionConnIds[sessionId] = append(sessionConnIds[sessionId], connId)
|
|
}
|
|
|
|
// Mark half of the connections connected, close the other half.
|
|
for i, connId := range connIds {
|
|
if i%2 == 0 {
|
|
_, cs, err := repo.ConnectConnection(ctx, ConnectWith{
|
|
ConnectionId: connId,
|
|
ClientTcpAddress: "127.0.0.1",
|
|
ClientTcpPort: 22,
|
|
EndpointTcpAddress: "127.0.0.1",
|
|
EndpointTcpPort: 22,
|
|
})
|
|
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 {
|
|
resp, err := repo.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)
|
|
}
|
|
}
|
|
|
|
// There is a 10 second delay to account for time for the connections to
|
|
// transition
|
|
time.Sleep(15 * time.Second)
|
|
|
|
// Now we try some scenarios.
|
|
{
|
|
// First test an empty set.
|
|
result, err := repo.ShouldCloseConnectionsOnWorker(ctx, nil, nil)
|
|
require.NoError(err)
|
|
require.Zero(result, "should be empty when no connections are supplied")
|
|
}
|
|
|
|
{
|
|
// Here we pass in all of our connections without a filter on
|
|
// session. This should return half of the connections - the ones
|
|
// that we marked as closed.
|
|
//
|
|
// Create a copy of our session map with the sessions that have
|
|
// closed connections taken out.
|
|
expectedSessionConnIds := make(map[string][]string)
|
|
for sessionId, connIds := range sessionConnIds {
|
|
for _, connId := range connIds {
|
|
if testIsConnectionClosed(ctx, t, repo, connId) {
|
|
expectedSessionConnIds[sessionId] = append(expectedSessionConnIds[sessionId], connId)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Send query, use all connections w/o a filter on sessions.
|
|
actualSessionConnIds, err := repo.ShouldCloseConnectionsOnWorker(ctx, connIds, nil)
|
|
require.NoError(err)
|
|
require.Equal(expectedSessionConnIds, actualSessionConnIds)
|
|
}
|
|
|
|
{
|
|
// Finally, add a session filter. We do this by just alternating
|
|
// the session IDs we want to filter on.
|
|
expectedSessionConnIds := make(map[string][]string)
|
|
var filterSessionIds []string
|
|
var filterSession bool
|
|
for sessionId, connIds := range sessionConnIds {
|
|
for _, connId := range connIds {
|
|
if testIsConnectionClosed(ctx, t, repo, connId) {
|
|
if !filterSession {
|
|
expectedSessionConnIds[sessionId] = append(expectedSessionConnIds[sessionId], connId)
|
|
} else {
|
|
filterSessionIds = append(filterSessionIds, sessionId)
|
|
}
|
|
|
|
// Toggle filterSession here (instead of just outer session
|
|
// loop) so that we aren't just lining up on
|
|
// connected/disconnected connections.
|
|
filterSession = !filterSession
|
|
}
|
|
}
|
|
}
|
|
|
|
// Send query with the session filter.
|
|
actualSessionConnIds, err := repo.ShouldCloseConnectionsOnWorker(ctx, connIds, filterSessionIds)
|
|
require.NoError(err)
|
|
require.Equal(expectedSessionConnIds, actualSessionConnIds)
|
|
}
|
|
}
|
|
|
|
func testIsConnectionClosed(ctx context.Context, t *testing.T, repo *Repository, connId string) bool {
|
|
require := require.New(t)
|
|
_, states, err := repo.LookupConnection(ctx, connId)
|
|
require.NoError(err)
|
|
// Use first state as this LookupConnections returns ordered by
|
|
// start time, descending.
|
|
return states[0].Status == StatusClosed
|
|
}
|