From 4de646497816733afba9c5f933ff9398eed5fc59 Mon Sep 17 00:00:00 2001
From: Timothy Messier <tim.messier@gmail.com>
Date: Fri, 4 Mar 2022 12:37:26 -0500
Subject: [PATCH] refact(session): Make CloseDeadConnectionsForWorker more
 testable

The query for CloseDeadConnectionsForWorker includes a time comparison
to exclude results that are too close to the current time. This is done
to account for a delay in reporting by the worker. However, since the
delay was hardcoded in the query, it made it difficult to test.

This moves the delay duration into a parameter that is set on the
repository. The amount of the delay can be set at repository creation
time using the new WithWorkerStateDelay option. This defaults to the
previously hardcoded amount.

This will allow tests to control the delay to exercise this
functionality fully.
---
 internal/session/options.go               | 15 ++++++++++++++-
 internal/session/query.go                 |  4 ++--
 internal/session/repository_connection.go | 22 ++++++++++++++--------
 3 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/internal/session/options.go b/internal/session/options.go
index 79fcac5bf..7d10ae818 100644
--- a/internal/session/options.go
+++ b/internal/session/options.go
@@ -1,6 +1,8 @@
 package session
 
 import (
+	"time"
+
 	"github.com/hashicorp/boundary/internal/db"
 	"github.com/hashicorp/boundary/internal/db/timestamp"
 )
@@ -29,10 +31,13 @@ type options struct {
 	withSessionIds        []string
 	withServerId          string
 	withDbOpts            []db.Option
+	withWorkerStateDelay  time.Duration
 }
 
 func getDefaultOptions() options {
-	return options{}
+	return options{
+		withWorkerStateDelay: 10 * time.Second,
+	}
 }
 
 // WithLimit provides an option to provide a limit. Intentionally allowing
@@ -108,3 +113,11 @@ func WithDbOpts(opts ...db.Option) Option {
 		o.withDbOpts = opts
 	}
 }
+
+// WithWorkerStateDelay is used byt queries to account for a delay in state
+// propagation between worker and controller.
+func WithWorkerStateDelay(d time.Duration) Option {
+	return func(o *options) {
+		o.withWorkerStateDelay = d
+	}
+}
diff --git a/internal/session/query.go b/internal/session/query.go
index a33e053bd..901b01f6d 100644
--- a/internal/session/query.go
+++ b/internal/session/query.go
@@ -316,14 +316,14 @@ with
       -- It's not in limbo between when it moved into this state and when
       -- it started being reported by the worker, which is roughly every
       -- 2-3 seconds
-      start_time < wt_sub_seconds_from_now(10)
+      start_time < wt_sub_seconds_from_now(@worker_state_delay_seconds)
   ),
   connections_to_close as (
     select public_id
       from session_connection
     where
       -- Related to the worker that just reported to us
-      server_id = ?
+      server_id = @server_id
         and
       -- These are connection IDs that just got reported to us by the given
       -- worker, so they should not be considered closed.
diff --git a/internal/session/repository_connection.go b/internal/session/repository_connection.go
index 3639da422..e63b0bd20 100644
--- a/internal/session/repository_connection.go
+++ b/internal/session/repository_connection.go
@@ -24,6 +24,10 @@ type ConnectionRepository struct {
 
 	// defaultLimit provides a default for limiting the number of results returned from the repo
 	defaultLimit int
+
+	// workerStateDelay is used by queries to account for a delay in state propagation between
+	// worker and controller
+	workerStateDelay time.Duration
 }
 
 // NewConnectionRepository creates a new session Connection Repository. Supports the options: WithLimit
@@ -45,10 +49,11 @@ func NewConnectionRepository(ctx context.Context, r db.Reader, w db.Writer, kms
 		opts.withLimit = db.DefaultLimit
 	}
 	return &ConnectionRepository{
-		reader:       r,
-		writer:       w,
-		kms:          kms,
-		defaultLimit: opts.withLimit,
+		reader:           r,
+		writer:           w,
+		kms:              kms,
+		defaultLimit:     opts.withLimit,
+		workerStateDelay: opts.withWorkerStateDelay,
 	}, nil
 }
 
@@ -359,16 +364,17 @@ func (r *ConnectionRepository) CloseDeadConnectionsForWorker(ctx context.Context
 		return db.NoRowsAffected, errors.New(ctx, errors.InvalidParameter, op, "missing server id")
 	}
 
-	args := make([]interface{}, 0, len(foundConns)+1)
-	args = append(args, serverId)
+	args := make([]interface{}, 0, len(foundConns)+2)
+	args = append(args, sql.Named("server_id", serverId))
+	args = append(args, sql.Named("worker_state_delay_seconds", r.workerStateDelay.Seconds()))
 
 	var publicIdStr string
 	if len(foundConns) > 0 {
 		publicIdStr = `public_id not in (%s) and`
 		params := make([]string, len(foundConns))
 		for i, connId := range foundConns {
-			params[i] = fmt.Sprintf("@%d", i+2) // Add one for server ID, and offsets start at 1
-			args = append(args, sql.Named(fmt.Sprintf("%d", i+2), connId))
+			params[i] = fmt.Sprintf("@%d", i)
+			args = append(args, sql.Named(fmt.Sprintf("%d", i), connId))
 		}
 		publicIdStr = fmt.Sprintf(publicIdStr, strings.Join(params, ","))
 	}