tstime: replace time.NewTimer with tstime.Clock timerclairew/refactor-new-timer

Updates #8587

Signed-off-by: Claire Wang <claire@tailscale.com>

6 files changed, 38 insertions, 28 deletions

diff --git a/control/controlclient/auto.go b/control/controlclient/auto.go
index dab37a510..7c143980d 100644
--- a/control/controlclient/auto.go
+++ b/control/controlclient/auto.go
@@ -15,6 +15,7 @@ import (
 	"tailscale.com/logtail/backoff"
 	"tailscale.com/net/sockstats"
 	"tailscale.com/tailcfg"
+	"tailscale.com/tstime"
 	"tailscale.com/types/empty"
 	"tailscale.com/types/key"
 	"tailscale.com/types/logger"
@@ -48,7 +49,7 @@ var _ Client = (*Auto)(nil)
 // It's a concrete implementation of the Client interface.
 type Auto struct {
 	direct     *Direct // our interface to the server APIs
-	timeNow    func() time.Time
+	clock      tstime.Clock
 	logf       logger.Logf
 	expiry     *time.Time
 	closed     bool
@@ -107,12 +108,12 @@ func NewNoStart(opts Options) (_ *Auto, err error) {
 	if opts.Logf == nil {
 		opts.Logf = func(fmt string, args ...any) {}
 	}
-	if opts.TimeNow == nil {
-		opts.TimeNow = time.Now
+	if opts.Clock == nil {
+		opts.Clock = &tstime.StdClock{}
 	}
 	c := &Auto{
 		direct:     direct,
-		timeNow:    opts.TimeNow,
+		clock:      opts.Clock,
 		logf:       opts.Logf,
 		newMapCh:   make(chan struct{}, 1),
 		quit:       make(chan struct{}),
@@ -702,14 +703,14 @@ func (c *Auto) Logout(ctx context.Context) error {
 	c.mu.Unlock()
 	c.cancelAuth()
 
-	timer := time.NewTimer(10 * time.Second)
+	timer, timerChannel := c.clock.NewTimer(10 * time.Second)
 	defer timer.Stop()
 	select {
 	case err := <-errc:
 		return err
 	case <-ctx.Done():
 		return ctx.Err()
-	case <-timer.C:
+	case <-timerChannel:
 		return context.DeadlineExceeded
 	}
 }
@@ -770,7 +771,7 @@ func (c *Auto) TestOnlySetAuthKey(authkey string) {
 }
 
 func (c *Auto) TestOnlyTimeNow() time.Time {
-	return c.timeNow()
+	return c.clock.Now()
 }
 
 // SetDNS sends the SetDNSRequest request to the control plane server,
diff --git a/control/controlclient/direct.go b/control/controlclient/direct.go
index 006f2614a..54d5ea895 100644
--- a/control/controlclient/direct.go
+++ b/control/controlclient/direct.go
@@ -45,6 +45,7 @@ import (
 	"tailscale.com/syncs"
 	"tailscale.com/tailcfg"
 	"tailscale.com/tka"
+	"tailscale.com/tstime"
 	"tailscale.com/types/key"
 	"tailscale.com/types/logger"
 	"tailscale.com/types/netmap"
@@ -63,7 +64,7 @@ type Direct struct {
 	dialer                 *tsdial.Dialer
 	dnsCache               *dnscache.Resolver
 	serverURL              string // URL of the tailcontrol server
-	timeNow                func() time.Time
+	clock                  tstime.Clock
 	lastPrintMap           time.Time
 	newDecompressor        func() (Decompressor, error)
 	keepAlive              bool
@@ -105,8 +106,8 @@ type Options struct {
 	GetMachinePrivateKey func() (key.MachinePrivate, error) // returns the machine key to use
 	ServerURL            string                             // URL of the tailcontrol server
 	AuthKey              string                             // optional node auth key for auto registration
-	TimeNow              func() time.Time                   // time.Now implementation used by Client
-	Hostinfo             *tailcfg.Hostinfo                  // non-nil passes ownership, nil means to use default using os.Hostname, etc
+	Clock                tstime.Clock
+	Hostinfo             *tailcfg.Hostinfo // non-nil passes ownership, nil means to use default using os.Hostname, etc
 	DiscoPublicKey       key.DiscoPublic
 	NewDecompressor      func() (Decompressor, error)
 	KeepAlive            bool
@@ -191,8 +192,8 @@ func NewDirect(opts Options) (*Direct, error) {
 	if err != nil {
 		return nil, err
 	}
-	if opts.TimeNow == nil {
-		opts.TimeNow = time.Now
+	if opts.Clock == nil {
+		opts.Clock = &tstime.StdClock{}
 	}
 	if opts.Logf == nil {
 		// TODO(apenwarr): remove this default and fail instead.
@@ -235,7 +236,7 @@ func NewDirect(opts Options) (*Direct, error) {
 		httpc:                  httpc,
 		getMachinePrivKey:      opts.GetMachinePrivateKey,
 		serverURL:              opts.ServerURL,
-		timeNow:                opts.TimeNow,
+		clock:                  opts.Clock,
 		logf:                   opts.Logf,
 		newDecompressor:        opts.NewDecompressor,
 		keepAlive:              opts.KeepAlive,
@@ -432,7 +433,7 @@ func (c *Direct) doLogin(ctx context.Context, opt loginOpt) (mustRegen bool, new
 	authKey, isWrapped, wrappedSig, wrappedKey := decodeWrappedAuthkey(c.authKey, c.logf)
 	hi := c.hostInfoLocked()
 	backendLogID := hi.BackendLogID
-	expired := c.expiry != nil && !c.expiry.IsZero() && c.expiry.Before(c.timeNow())
+	expired := c.expiry != nil && !c.expiry.IsZero() && c.expiry.Before(c.clock.Now())
 	c.mu.Unlock()
 
 	machinePrivKey, err := c.getMachinePrivKey()
@@ -947,7 +948,7 @@ func (c *Direct) sendMapRequest(ctx context.Context, maxPolls int, readOnly bool
 		return nil
 	}
 
-	timeout := time.NewTimer(pollTimeout)
+	timeout, timeoutChannel := c.clock.NewTimer(pollTimeout)
 	timeoutReset := make(chan struct{})
 	pollDone := make(chan struct{})
 	defer close(pollDone)
@@ -957,14 +958,14 @@ func (c *Direct) sendMapRequest(ctx context.Context, maxPolls int, readOnly bool
 			case <-pollDone:
 				vlogf("netmap: ending timeout goroutine")
 				return
-			case <-timeout.C:
+			case <-timeoutChannel:
 				c.logf("map response long-poll timed out!")
 				cancel()
 				return
 			case <-timeoutReset:
 				if !timeout.Stop() {
 					select {
-					case <-timeout.C:
+					case <-timeoutChannel:
 					case <-pollDone:
 						vlogf("netmap: ending timeout goroutine")
 						return
@@ -1089,7 +1090,7 @@ func (c *Direct) sendMapRequest(ctx context.Context, maxPolls int, readOnly bool
 				go dumpGoroutinesToURL(c.httpc, resp.Debug.GoroutineDumpURL)
 			}
 			if sleep := time.Duration(resp.Debug.SleepSeconds * float64(time.Second)); sleep > 0 {
-				if err := sleepAsRequested(ctx, c.logf, timeoutReset, sleep); err != nil {
+				if err := sleepAsRequested(ctx, c.logf, timeoutReset, sleep, c.clock); err != nil {
 					return err
 				}
 			}
@@ -1119,7 +1120,7 @@ func (c *Direct) sendMapRequest(ctx context.Context, maxPolls int, readOnly bool
 		// This is handy for debugging, and our logs processing
 		// pipeline depends on it. (TODO: Remove this dependency.)
 		// Code elsewhere prints netmap diffs every time they are received.
-		now := c.timeNow()
+		now := c.clock.Now()
 		if now.Sub(c.lastPrintMap) >= 5*time.Minute {
 			c.lastPrintMap = now
 			c.logf("[v1] new network map[%d]:\n%s", i, nm.VeryConcise())
@@ -1459,7 +1460,7 @@ func answerC2NPing(logf logger.Logf, c2nHandler http.Handler, c *http.Client, pr
 	}
 }
 
-func sleepAsRequested(ctx context.Context, logf logger.Logf, timeoutReset chan<- struct{}, d time.Duration) error {
+func sleepAsRequested(ctx context.Context, logf logger.Logf, timeoutReset chan<- struct{}, d time.Duration, clock tstime.Clock) error {
 	const maxSleep = 5 * time.Minute
 	if d > maxSleep {
 		logf("sleeping for %v, capped from server-requested %v ...", maxSleep, d)
@@ -1468,20 +1469,20 @@ func sleepAsRequested(ctx context.Context, logf logger.Logf, timeoutReset chan<-
 		logf("sleeping for server-requested %v ...", d)
 	}
 
-	ticker := time.NewTicker(pollTimeout / 2)
+	ticker, tickerChannel := clock.NewTicker(pollTimeout / 2)
 	defer ticker.Stop()
-	timer := time.NewTimer(d)
+	timer, timerChannel := clock.NewTimer(d)
 	defer timer.Stop()
 	for {
 		select {
 		case <-ctx.Done():
 			return ctx.Err()
-		case <-timer.C:
+		case <-timerChannel:
 			return nil
-		case <-ticker.C:
+		case <-tickerChannel:
 			select {
 			case timeoutReset <- struct{}{}:
-			case <-timer.C:
+			case <-timerChannel:
 				return nil
 			case <-ctx.Done():
 				return ctx.Err()
diff --git a/control/controlclient/noise.go b/control/controlclient/noise.go
index 4a8335133..e3a8a9b9e 100644
--- a/control/controlclient/noise.go
+++ b/control/controlclient/noise.go
@@ -23,6 +23,7 @@ import (
 	"tailscale.com/net/netmon"
 	"tailscale.com/net/tsdial"
 	"tailscale.com/tailcfg"
+	"tailscale.com/tstime"
 	"tailscale.com/types/key"
 	"tailscale.com/types/logger"
 	"tailscale.com/util/mak"
@@ -450,6 +451,7 @@ func (nc *NoiseClient) dial(ctx context.Context) (*noiseConn, error) {
 		DialPlan:        dialPlan,
 		Logf:            nc.logf,
 		NetMon:          nc.netMon,
+		Clock:           &tstime.StdClock{},
 	}).Dial(ctx)
 	if err != nil {
 		return nil, err
diff --git a/control/controlhttp/client.go b/control/controlhttp/client.go
index b0d91bada..facd0f1b1 100644
--- a/control/controlhttp/client.go
+++ b/control/controlhttp/client.go
@@ -147,13 +147,13 @@ func (a *Dialer) dial(ctx context.Context) (*ClientConn, error) {
 			// before we do anything.
 			if c.DialStartDelaySec > 0 {
 				a.logf("[v2] controlhttp: waiting %.2f seconds before dialing %q @ %v", c.DialStartDelaySec, a.Hostname, c.IP)
-				tmr := time.NewTimer(time.Duration(c.DialStartDelaySec * float64(time.Second)))
+				tmr, tmrChannel := a.Clock.NewTimer(time.Duration(c.DialStartDelaySec * float64(time.Second)))
 				defer tmr.Stop()
 				select {
 				case <-ctx.Done():
 					err = ctx.Err()
 					return
-				case <-tmr.C:
+				case <-tmrChannel:
 				}
 			}
 
@@ -319,7 +319,7 @@ func (a *Dialer) dialHost(ctx context.Context, addr netip.Addr) (*ClientConn, er
 
 	// In case outbound port 80 blocked or MITM'ed poorly, start a backup timer
 	// to dial port 443 if port 80 doesn't either succeed or fail quickly.
-	try443Timer := time.AfterFunc(a.httpsFallbackDelay(), func() { try(u443) })
+	try443Timer := a.Clock.AfterFunc(a.httpsFallbackDelay(), func() { try(u443) })
 	defer try443Timer.Stop()
 
 	var err80, err443 error
diff --git a/control/controlhttp/constants.go b/control/controlhttp/constants.go
index b838f84c4..76c76699f 100644
--- a/control/controlhttp/constants.go
+++ b/control/controlhttp/constants.go
@@ -11,6 +11,7 @@ import (
 	"tailscale.com/net/dnscache"
 	"tailscale.com/net/netmon"
 	"tailscale.com/tailcfg"
+	"tailscale.com/tstime"
 	"tailscale.com/types/key"
 	"tailscale.com/types/logger"
 )
@@ -89,6 +90,8 @@ type Dialer struct {
 	drainFinished        chan struct{}
 	omitCertErrorLogging bool
 	testFallbackDelay    time.Duration
+
+	Clock tstime.Clock
 }
 
 func strDef(v1, v2 string) string {
diff --git a/control/controlhttp/http_test.go b/control/controlhttp/http_test.go
index 3cf790d8c..c5fcdf5c3 100644
--- a/control/controlhttp/http_test.go
+++ b/control/controlhttp/http_test.go
@@ -25,6 +25,7 @@ import (
 	"tailscale.com/net/socks5"
 	"tailscale.com/net/tsdial"
 	"tailscale.com/tailcfg"
+	"tailscale.com/tstest"
 	"tailscale.com/types/key"
 	"tailscale.com/types/logger"
 )
@@ -204,6 +205,7 @@ func testControlHTTP(t *testing.T, param httpTestParam) {
 		Logf:                 t.Logf,
 		omitCertErrorLogging: true,
 		testFallbackDelay:    50 * time.Millisecond,
+		Clock:                &tstest.Clock{},
 	}
 
 	if proxy != nil {
@@ -660,6 +662,7 @@ func TestDialPlan(t *testing.T) {
 				drainFinished:        drained,
 				omitCertErrorLogging: true,
 				testFallbackDelay:    50 * time.Millisecond,
+				Clock:                &tstest.Clock{},
 			}
 
 			conn, err := a.dial(ctx)