Add Infra

2025-03-03 18:59:25 +07:00
parent c02d8cda69
commit c9f1e33c24
1782 changed files with 692894 additions and 21 deletions
@@ -0,0 +1,260 @@
+// Copyright (c) 2024 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package fxclock
+
+import (
+	"context"
+	"sort"
+	"sync"
+	"time"
+)
+
+// Clock defines how Fx accesses time.
+// We keep the interface pretty minimal.
+type Clock interface {
+	Now() time.Time
+	Since(time.Time) time.Duration
+	Sleep(time.Duration)
+	WithTimeout(context.Context, time.Duration) (context.Context, context.CancelFunc)
+}
+
+// System is the default implementation of Clock based on real time.
+var System Clock = systemClock{}
+
+type systemClock struct{}
+
+func (systemClock) Now() time.Time {
+	return time.Now()
+}
+
+func (systemClock) Since(t time.Time) time.Duration {
+	return time.Since(t)
+}
+
+func (systemClock) Sleep(d time.Duration) {
+	time.Sleep(d)
+}
+
+func (systemClock) WithTimeout(ctx context.Context, d time.Duration) (context.Context, context.CancelFunc) {
+	return context.WithTimeout(ctx, d)
+}
+
+// Mock adapted from
+// https://github.com/uber-go/zap/blob/7db06bc9b095571d3dc3d4eebdfbe4dd9bd20405/internal/ztest/clock.go.
+
+// Mock is a fake source of time.
+// It implements standard time operations,
+// but allows the user to control the passage of time.
+//
+// Use the [Add] method to progress time.
+type Mock struct {
+	mu  sync.RWMutex
+	now time.Time
+
+	// The MockClock works by maintaining a list of waiters.
+	// Each waiter knows the time at which it should be resolved.
+	// When the clock advances, all waiters that are in range are resolved
+	// in chronological order.
+	waiters     []waiter
+	waiterAdded *sync.Cond
+}
+
+var _ Clock = (*Mock)(nil)
+
+// NewMock builds a new mock clock
+// using the current actual time as the initial time.
+func NewMock() *Mock {
+	m := &Mock{now: time.Now()}
+	m.waiterAdded = sync.NewCond(&m.mu)
+	return m
+}
+
+// Now reports the current time.
+func (c *Mock) Now() time.Time {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	return c.now
+}
+
+// Since reports the time elapsed since t.
+// This is short for Now().Sub(t).
+func (c *Mock) Since(t time.Time) time.Duration {
+	return c.Now().Sub(t)
+}
+
+// Sleep pauses the current goroutine for the given duration.
+//
+// With the mock clock, this will freeze
+// until the clock is advanced with [Add] past the deadline.
+func (c *Mock) Sleep(d time.Duration) {
+	ch := make(chan struct{})
+	c.runAt(c.Now().Add(d), func() { close(ch) })
+	<-ch
+}
+
+// WithTimeout returns a new context with a deadline of now + d.
+//
+// When the deadline is passed, the returned context's Done channel is closed
+// and the context's Err method returns context.DeadlineExceeded.
+// If the cancel function is called before the deadline is passed,
+// the context's Err method returns context.Canceled.
+func (c *Mock) WithTimeout(ctx context.Context, d time.Duration) (context.Context, context.CancelFunc) {
+	// Unfortunately, we can't use context.WithCancelCause here.
+	// Per its documentation (and verified by trying it):
+	//
+	//   ctx, cancel := context.WithCancelCause(parent)
+	//   cancel(myError)
+	//   ctx.Err() // returns context.Canceled
+	//   context.Cause(ctx) // returns myError
+	//
+	// So it won't do for our purposes.
+	deadline := c.Now().Add(d)
+	inner, cancelInner := context.WithCancel(ctx)
+	dctx := &deadlineCtx{
+		inner:       inner,
+		cancelInner: cancelInner,
+		done:        make(chan struct{}),
+		deadline:    deadline,
+	}
+	ctx = dctx
+
+	c.runAt(deadline, func() {
+		dctx.cancel(context.DeadlineExceeded)
+	})
+	return ctx, func() { dctx.cancel(context.Canceled) }
+}
+
+type deadlineCtx struct {
+	inner       context.Context
+	cancelInner func()
+
+	done     chan struct{}
+	deadline time.Time
+
+	mu  sync.Mutex // guards err; the rest is immutable
+	err error
+}
+
+var _ context.Context = (*deadlineCtx)(nil)
+
+func (c *deadlineCtx) Deadline() (deadline time.Time, ok bool) { return c.deadline, true }
+func (c *deadlineCtx) Done() <-chan struct{}                   { return c.done }
+func (c *deadlineCtx) Value(key any) any                       { return c.inner.Value(key) }
+
+func (c *deadlineCtx) Err() error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.err
+}
+
+func (c *deadlineCtx) cancel(err error) {
+	c.mu.Lock()
+	if c.err == nil {
+		c.err = err
+		close(c.done)
+		c.cancelInner()
+	}
+	c.mu.Unlock()
+}
+
+// runAt schedules the given function to be run at the given time.
+// The function runs without a lock held, so it may schedule more work.
+func (c *Mock) runAt(t time.Time, fn func()) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	c.waiters = append(c.waiters, waiter{until: t, fn: fn})
+	c.waiterAdded.Broadcast()
+}
+
+// AwaitScheduled blocks until there are at least N
+// operations scheduled for the future.
+func (c *Mock) AwaitScheduled(n int) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	// Note: waiterAdded is associated with c.mu,
+	// the same lock we're holding here.
+	//
+	// When we call Wait(), it'll release the lock
+	// and block until signaled by runAt,
+	// at which point it'll reacquire the lock
+	// (waiting until runAt has released it).
+	for len(c.waiters) < n {
+		c.waiterAdded.Wait()
+	}
+}
+
+type waiter struct {
+	until time.Time
+	fn    func()
+}
+
+// Add progresses time by the given duration.
+// Other operations waiting for the time to advance
+// will be resolved if they are within range.
+//
+// Side effects of operations waiting for the time to advance
+// will take effect on a best-effort basis.
+// Avoid racing with operations that have side effects.
+//
+// Panics if the duration is negative.
+func (c *Mock) Add(d time.Duration) {
+	if d < 0 {
+		panic("cannot add negative duration")
+	}
+
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	sort.Slice(c.waiters, func(i, j int) bool {
+		return c.waiters[i].until.Before(c.waiters[j].until)
+	})
+
+	newTime := c.now.Add(d)
+	// newTime won't be recorded until the end of this method.
+	// This ensures that any waiters that are resolved
+	// are resolved at the time they were expecting.
+
+	for len(c.waiters) > 0 {
+		w := c.waiters[0]
+		if w.until.After(newTime) {
+			break
+		}
+		c.waiters[0] = waiter{} // avoid memory leak
+		c.waiters = c.waiters[1:]
+
+		// The waiter is within range.
+		// Travel to the time of the waiter and resolve it.
+		c.now = w.until
+
+		// The waiter may schedule more work
+		// so we must release the lock.
+		c.mu.Unlock()
+		w.fn()
+		// Sleeping here is necessary to let the side effects of waiters
+		// take effect before we continue.
+		time.Sleep(1 * time.Millisecond)
+		c.mu.Lock()
+	}
+
+	c.now = newTime
+}