bootstrap: Use ticker to get rid of test timeouts

The bootstrapping tests still need to use timeouts because of the grace
period, which is the duration we give the server to exit before we force
the upgrade. Convert the interface to instead accept tickers such that
we can get rid of using timeouts in our tests.

2 files changed, 30 insertions, 17 deletions

diff --git a/internal/bootstrap/bootstrap.go b/internal/bootstrap/bootstrap.go
index f1d355926..3cb08d6fa 100644
--- a/internal/bootstrap/bootstrap.go
+++ b/internal/bootstrap/bootstrap.go
@@ -6,10 +6,10 @@ import (
 	"os"
 	"os/signal"
 	"syscall"
-	"time"
 
 	"github.com/cloudflare/tableflip"
 	log "github.com/sirupsen/logrus"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/helper"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/helper/env"
 	"golang.org/x/sys/unix"
 )
@@ -29,7 +29,7 @@ type Listener interface {
 	// Start starts all registered starters to accept connections.
 	Start() error
 	// Wait terminates all registered starters.
-	Wait(gracefulTimeout time.Duration, stopAction func()) error
+	Wait(gracePeriodTicker helper.Ticker, stopAction func()) error
 }
 
 // Bootstrap handles graceful upgrades
@@ -160,7 +160,7 @@ func (b *Bootstrap) Start() error {
 // SIGTERM, SIGINT and a runtime error will trigger an immediate shutdown
 // in case of an upgrade there will be a grace period to complete the ongoing requests
 // stopAction will be invoked during a graceful stop. It must wait until the shutdown is completed.
-func (b *Bootstrap) Wait(gracefulTimeout time.Duration, stopAction func()) error {
+func (b *Bootstrap) Wait(gracePeriodTicker helper.Ticker, stopAction func()) error {
 	signals := []os.Signal{syscall.SIGTERM, syscall.SIGINT}
 	immediateShutdown := make(chan os.Signal, len(signals))
 	signal.Notify(immediateShutdown, signals...)
@@ -176,7 +176,7 @@ func (b *Bootstrap) Wait(gracefulTimeout time.Duration, stopAction func()) error
 		// the new process signaled its readiness and we started a graceful stop
 		// however no further upgrades can be started until this process is running
 		// we set a grace period and then we force a termination.
-		waitError := b.waitGracePeriod(gracefulTimeout, immediateShutdown, stopAction)
+		waitError := b.waitGracePeriod(gracePeriodTicker, immediateShutdown, stopAction)
 
 		err = fmt.Errorf("graceful upgrade: %v", waitError)
 	case s := <-immediateShutdown:
@@ -188,8 +188,8 @@ func (b *Bootstrap) Wait(gracefulTimeout time.Duration, stopAction func()) error
 	return err
 }
 
-func (b *Bootstrap) waitGracePeriod(gracefulTimeout time.Duration, kill <-chan os.Signal, stopAction func()) error {
-	log.WithField("graceful_timeout", gracefulTimeout).Warn("starting grace period")
+func (b *Bootstrap) waitGracePeriod(gracePeriodTicker helper.Ticker, kill <-chan os.Signal, stopAction func()) error {
+	log.Warn("starting grace period")
 
 	allServersDone := make(chan struct{})
 	go func() {
@@ -199,8 +199,10 @@ func (b *Bootstrap) waitGracePeriod(gracefulTimeout time.Duration, kill <-chan o
 		close(allServersDone)
 	}()
 
+	gracePeriodTicker.Reset()
+
 	select {
-	case <-time.After(gracefulTimeout):
+	case <-gracePeriodTicker.C():
 		return fmt.Errorf("grace period expired")
 	case <-kill:
 		return fmt.Errorf("force shutdown")
@@ -249,7 +251,7 @@ func (n *Noop) Start() error {
 }
 
 // Wait terminates all registered starters.
-func (n *Noop) Wait(_ time.Duration, stopAction func()) error {
+func (n *Noop) Wait(_ helper.Ticker, stopAction func()) error {
 	select {
 	case <-n.shutdown:
 		if stopAction != nil {
diff --git a/internal/bootstrap/bootstrap_test.go b/internal/bootstrap/bootstrap_test.go
index 7491002cd..872a62ec0 100644
--- a/internal/bootstrap/bootstrap_test.go
+++ b/internal/bootstrap/bootstrap_test.go
@@ -8,10 +8,10 @@ import (
 	"path/filepath"
 	"syscall"
 	"testing"
-	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/helper"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/testhelper"
 )
 
@@ -93,7 +93,7 @@ func TestBootstrap_listenerError(t *testing.T) {
 	b, upgrader, listeners := setup(t, ctx)
 
 	waitCh := make(chan error)
-	go func() { waitCh <- b.Wait(time.Hour, nil) }()
+	go func() { waitCh <- b.Wait(helper.NewManualTicker(), nil) }()
 
 	// Signal readiness, but don't start the upgrade. Like this, we can close the listener in a
 	// raceless manner and wait for the error to propagate.
@@ -114,7 +114,7 @@ func TestBootstrap_signal(t *testing.T) {
 			b, upgrader, _ := setup(t, ctx)
 
 			waitCh := make(chan error)
-			go func() { waitCh <- b.Wait(time.Hour, nil) }()
+			go func() { waitCh <- b.Wait(helper.NewManualTicker(), nil) }()
 
 			// Start the upgrade, but don't unblock `Exit()` such that we'll be blocked
 			// waiting on the parent.
@@ -137,10 +137,14 @@ func TestBootstrap_gracefulTerminationStuck(t *testing.T) {
 
 	b, upgrader, _ := setup(t, ctx)
 
+	gracePeriodTicker := helper.NewManualTicker()
+
 	doneCh := make(chan struct{})
-	err := performUpgrade(t, b, upgrader, time.Second, nil, func() {
+	err := performUpgrade(t, b, upgrader, gracePeriodTicker, nil, func() {
 		defer close(doneCh)
 
+		gracePeriodTicker.Tick()
+
 		// We block on context cancellation here, which essentially means that this won't
 		// terminate and thus the graceful termination will be stuck.
 		<-ctx.Done()
@@ -160,7 +164,7 @@ func TestBootstrap_gracefulTerminationWithSignals(t *testing.T) {
 			b, upgrader, _ := setup(t, ctx)
 
 			doneCh := make(chan struct{})
-			err := performUpgrade(t, b, upgrader, time.Hour, func() {
+			err := performUpgrade(t, b, upgrader, helper.NewManualTicker(), func() {
 				self, err := os.FindProcess(os.Getpid())
 				require.NoError(t, err)
 				require.NoError(t, self.Signal(sig))
@@ -184,14 +188,18 @@ func TestBootstrap_gracefulTerminationTimeoutWithListenerError(t *testing.T) {
 
 	b, upgrader, listeners := setup(t, ctx)
 
+	gracePeriodTicker := helper.NewManualTicker()
+
 	doneCh := make(chan struct{})
-	err := performUpgrade(t, b, upgrader, time.Second, nil, func() {
+	err := performUpgrade(t, b, upgrader, gracePeriodTicker, nil, func() {
 		defer close(doneCh)
 
 		// We inject an error into the Unix socket to assert that this won't kill the server
 		// immediately, but waits for the TCP connection to terminate as expected.
 		listeners["unix"].errorCh <- assert.AnError
 
+		gracePeriodTicker.Tick()
+
 		// We block on context cancellation here, which essentially means that this won't
 		// terminate.
 		<-ctx.Done()
@@ -208,7 +216,10 @@ func TestBootstrap_gracefulTermination(t *testing.T) {
 
 	b, upgrader, _ := setup(t, ctx)
 
-	require.Equal(t, fmt.Errorf("graceful upgrade: completed"), performUpgrade(t, b, upgrader, time.Hour, nil, nil))
+	require.Equal(t,
+		fmt.Errorf("graceful upgrade: completed"),
+		performUpgrade(t, b, upgrader, helper.NewManualTicker(), nil, nil),
+	)
 }
 
 func TestBootstrap_portReuse(t *testing.T) {
@@ -233,12 +244,12 @@ func performUpgrade(
 	t *testing.T,
 	b *Bootstrap,
 	upgrader *mockUpgrader,
-	gracefulWait time.Duration,
+	gracePeriodTicker helper.Ticker,
 	duringGracePeriodCallback func(),
 	stopAction func(),
 ) error {
 	waitCh := make(chan error)
-	go func() { waitCh <- b.Wait(gracefulWait, stopAction) }()
+	go func() { waitCh <- b.Wait(gracePeriodTicker, stopAction) }()
 
 	// Simulate an upgrade request after entering into the blocking b.Wait() and during the
 	// slowRequest execution