tools/lib/clock/clock.go

// Copyright 2021 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package clock

import (
	"context"
	"sync"
	"time"
)

type clock interface {
	Now() time.Time
	After(d time.Duration) <-chan time.Time
	Sleep(d time.Duration)
}

type clockKeyType string

// clockKey is the key we use to associate a clock with a Context.
const clockKey = clockKeyType("clock")

// Now returns the current time for the clock associated with the given context,
// or the real current time if there is no clock associated with the context.
// Any code that needs to be tested with a mocked-out time should use
// `clock.Now()` instead of `time.Now()`.
func Now(ctx context.Context) time.Time {
	if c, ok := ctx.Value(clockKey).(clock); ok && c != nil {
		return c.Now()
	}
	return time.Now()
}

// After returns time.After() or the equivalent for the clock associated with
// the given context.
func After(ctx context.Context, d time.Duration) <-chan time.Time {
	if c, ok := ctx.Value(clockKey).(clock); ok && c != nil {
		return c.After(d)
	}
	return time.After(d)
}

// Sleep runs time.Sleep() or the equivalent for the clock associated with the
// given context.
func Sleep(ctx context.Context, d time.Duration) {
	if c, ok := ctx.Value(clockKey).(clock); ok && c != nil {
		c.Sleep(d)
	} else {
		time.Sleep(d)
	}
}

// NewContext returns a new context with the given clock attached.
//
// This should generally only be used in tests; production code should always
// use the real time.
func NewContext(ctx context.Context, c clock) context.Context {
	return context.WithValue(ctx, clockKey, c)
}

type timer struct {
	endTime time.Time
	ch      chan time.Time
}

// FakeClock provides support for mocking the current time in tests.
type FakeClock struct {
	mu            sync.Mutex
	now           time.Time
	pendingTimers []*timer
	afterCalled   chan struct{}
}

func NewFakeClock() *FakeClock {
	return &FakeClock{now: time.Now(), afterCalled: make(chan struct{})}
}

func (c *FakeClock) Now() time.Time {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.now
}

func (c *FakeClock) After(d time.Duration) <-chan time.Time {
	c.mu.Lock()
	defer c.mu.Unlock()
	t := &timer{c.now.Add(d), make(chan time.Time, 1)}
	c.pendingTimers = append(c.pendingTimers, t)
	select {
	case <-c.afterCalled: // Channel already closed.
	default:
		close(c.afterCalled)
	}
	return t.ch
}

func (c *FakeClock) Advance(d time.Duration) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.now = c.now.Add(d)

	// Notify timers that the time has changed.
	var pendingTimers []*timer
	for _, t := range c.pendingTimers {
		if !c.now.Before(t.endTime) {
			t.ch <- c.now
		} else {
			pendingTimers = append(pendingTimers, t)
		}
	}
	c.pendingTimers = pendingTimers
}

// AfterCalledChan returns the channel to wait for the clock's timer to be set from a
// call to After().
func (c *FakeClock) AfterCalledChan() <-chan struct{} {
	return c.afterCalled
}

func (c *FakeClock) Sleep(d time.Duration) {
	c.Advance(d)
}