From 20a6b7ce1afa89c90368913000645b758ed226a7 Mon Sep 17 00:00:00 2001
From: Taiwon Chung <taiwon@uber.com>
Date: Thu, 27 Jun 2024 15:44:25 -0400
Subject: [PATCH] feat: add function Runtime to dig.CallbackInfo
---
callback.go | 6 +++
constructor.go | 6 ++-
container.go | 19 ++++++++
decorate.go | 6 ++-
dig_int_test.go | 10 +++-
dig_test.go | 50 ++++++++++++++++++++
internal/digclock/clock.go | 83 +++++++++++++++++++++++++++++++++
internal/digclock/clock_test.go | 53 +++++++++++++++++++++
scope.go | 10 ++++
9 files changed, 238 insertions(+), 5 deletions(-)
create mode 100644 internal/digclock/clock.go
create mode 100644 internal/digclock/clock_test.go
diff --git a/callback.go b/callback.go
index dfe47ea6..6abc57d8 100644
--- a/callback.go
+++ b/callback.go
@@ -20,6 +20,8 @@
package dig
+import "time"
+
// CallbackInfo contains information about a provided function or decorator
// called by Dig, and is passed to a [Callback] registered with
// [WithProviderCallback] or [WithDecoratorCallback].
@@ -32,6 +34,10 @@ type CallbackInfo struct {
// function, if any. When used in conjunction with [RecoverFromPanics],
// this will be set to a [PanicError] when the function panics.
Error error
+
+ // Runtime contains the duration of time it took for the associated
+ // function to run.
+ Runtime time.Duration
}
// Callback is a function that can be registered with a provided function
diff --git a/constructor.go b/constructor.go
index 034c41c2..adec5fd5 100644
--- a/constructor.go
+++ b/constructor.go
@@ -161,11 +161,13 @@ func (n *constructorNode) Call(c containerStore) (err error) {
}
if n.callback != nil {
+ start := c.clock().Now()
// Wrap in separate func to include PanicErrors
defer func() {
n.callback(CallbackInfo{
- Name: fmt.Sprintf("%v.%v", n.location.Package, n.location.Name),
- Error: err,
+ Name: fmt.Sprintf("%v.%v", n.location.Package, n.location.Name),
+ Error: err,
+ Runtime: c.clock().Since(start),
})
}()
}
diff --git a/container.go b/container.go
index 983fd3f9..a875b5e0 100644
--- a/container.go
+++ b/container.go
@@ -25,6 +25,7 @@ import (
"math/rand"
"reflect"
+ "go.uber.org/dig/internal/digclock"
"go.uber.org/dig/internal/dot"
)
@@ -141,6 +142,9 @@ type containerStore interface {
// Returns invokerFn function to use when calling arguments.
invoker() invokerFn
+
+ // Returns a clock to use
+ clock() digclock.Clock
}
// New constructs a Container.
@@ -211,6 +215,21 @@ func (o setRandOption) applyOption(c *Container) {
c.scope.rand = o.r
}
+// Changes the source of time for the container.
+func setClock(c digclock.Clock) Option {
+ return setClockOption{c: c}
+}
+
+type setClockOption struct{ c digclock.Clock }
+
+func (o setClockOption) String() string {
+ return fmt.Sprintf("setClock(%v)", o.c)
+}
+
+func (o setClockOption) applyOption(c *Container) {
+ c.scope.clockSrc = o.c
+}
+
// DryRun is an Option which, when set to true, disables invocation of functions supplied to
// Provide and Invoke. Use this to build no-op containers.
func DryRun(dry bool) Option {
diff --git a/decorate.go b/decorate.go
index df362e98..f4c6be18 100644
--- a/decorate.go
+++ b/decorate.go
@@ -122,11 +122,13 @@ func (n *decoratorNode) Call(s containerStore) (err error) {
}
if n.callback != nil {
+ start := s.clock().Now()
// Wrap in separate func to include PanicErrors
defer func() {
n.callback(CallbackInfo{
- Name: fmt.Sprintf("%v.%v", n.location.Package, n.location.Name),
- Error: err,
+ Name: fmt.Sprintf("%v.%v", n.location.Package, n.location.Name),
+ Error: err,
+ Runtime: s.clock().Since(start),
})
}()
}
diff --git a/dig_int_test.go b/dig_int_test.go
index 47756076..89140816 100644
--- a/dig_int_test.go
+++ b/dig_int_test.go
@@ -20,8 +20,16 @@
package dig
-import "math/rand"
+import (
+ "math/rand"
+
+ "go.uber.org/dig/internal/digclock"
+)
func SetRand(r *rand.Rand) Option {
return setRand(r)
}
+
+func SetClock(c digclock.Clock) Option {
+ return setClock(c)
+}
diff --git a/dig_test.go b/dig_test.go
index 5cbf4ee8..e2f5bc67 100644
--- a/dig_test.go
+++ b/dig_test.go
@@ -34,6 +34,7 @@ import (
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.uber.org/dig"
+ "go.uber.org/dig/internal/digclock"
"go.uber.org/dig/internal/digtest"
)
@@ -1796,6 +1797,55 @@ func TestCallback(t *testing.T) {
})
}
+func TestCallbackRuntime(t *testing.T) {
+ t.Run("provided ctor runtime", func(t *testing.T) {
+ var called bool
+
+ mockClock := digclock.NewMock()
+ c := digtest.New(t, dig.SetClock(mockClock))
+ c.RequireProvide(
+ func() int {
+ mockClock.Add(1 * time.Millisecond)
+ return 5
+ },
+ dig.WithProviderCallback(func(ci dig.CallbackInfo) {
+ assert.Equal(t, "go.uber.org/dig_test.TestCallbackRuntime.func1.1", ci.Name)
+ assert.NoError(t, ci.Error)
+ assert.Equal(t, ci.Runtime, 1*time.Millisecond)
+
+ called = true
+ }),
+ )
+
+ c.Invoke(func(int) {})
+ assert.True(t, called)
+ })
+
+ t.Run("decorator runtime", func(t *testing.T) {
+ var called bool
+
+ mockClock := digclock.NewMock()
+ c := digtest.New(t, dig.SetClock(mockClock))
+ c.RequireProvide(giveInt)
+ c.RequireDecorate(
+ func(int) int {
+ mockClock.Add(1 * time.Millisecond)
+ return 10
+ },
+ dig.WithDecoratorCallback(func(ci dig.CallbackInfo) {
+ assert.Equal(t, "go.uber.org/dig_test.TestCallbackRuntime.func2.1", ci.Name)
+ assert.NoError(t, ci.Error)
+ assert.Equal(t, ci.Runtime, 1*time.Millisecond)
+
+ called = true
+ }),
+ )
+
+ c.Invoke(func(int) {})
+ assert.True(t, called)
+ })
+}
+
func TestProvideConstructorErrors(t *testing.T) {
t.Run("multiple-type constructor returns multiple objects of same type", func(t *testing.T) {
c := digtest.New(t)
diff --git a/internal/digclock/clock.go b/internal/digclock/clock.go
new file mode 100644
index 00000000..8c7d2d10
--- /dev/null
+++ b/internal/digclock/clock.go
@@ -0,0 +1,83 @@
+// 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 digclock
+
+import (
+ "time"
+)
+
+// Clock defines how dig accesses time.
+// We keep the interface pretty minimal.
+type Clock interface {
+ Now() time.Time
+ Since(time.Time) time.Duration
+}
+
+// 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)
+}
+
+// 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.
+//
+// Note that this implementation is not safe for concurrent use.
+type Mock struct {
+ now time.Time
+}
+
+var _ Clock = (*Mock)(nil)
+
+// NewMock creates a new mock clock with the current time set to the current time.
+func NewMock() *Mock {
+ return &Mock{now: time.Now()}
+}
+
+// Now returns the current time.
+func (m *Mock) Now() time.Time {
+ return m.now
+}
+
+// Since returns the time elapsed since the given time.
+func (m *Mock) Since(t time.Time) time.Duration {
+ return m.Now().Sub(t)
+}
+
+// Add progresses time by the given duration.
+//
+// It panics if the duration is negative.
+func (m *Mock) Add(d time.Duration) {
+ if d < 0 {
+ panic("cannot add negative duration")
+ }
+ m.now = m.now.Add(d)
+}
diff --git a/internal/digclock/clock_test.go b/internal/digclock/clock_test.go
new file mode 100644
index 00000000..fda79947
--- /dev/null
+++ b/internal/digclock/clock_test.go
@@ -0,0 +1,53 @@
+// 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 digclock
+
+import (
+ "testing"
+ "time"
+
+ "github.com/stretchr/testify/assert"
+)
+
+func TestSystemClock(t *testing.T) {
+ clock := System
+ testClock(t, clock, func(d time.Duration) { time.Sleep(d) })
+}
+
+func TestMockClock(t *testing.T) {
+ clock := NewMock()
+ testClock(t, clock, clock.Add)
+}
+
+func testClock(t *testing.T, clock Clock, advance func(d time.Duration)) {
+ now := clock.Now()
+ assert.False(t, now.IsZero())
+
+ t.Run("Since", func(t *testing.T) {
+ advance(1 * time.Millisecond)
+ assert.NotZero(t, clock.Since(now), "time must have advanced")
+ })
+}
+
+func TestMock_AddNegative(t *testing.T) {
+ clock := NewMock()
+ assert.Panics(t, func() { clock.Add(-1) })
+}
diff --git a/scope.go b/scope.go
index d5478aca..49b1a691 100644
--- a/scope.go
+++ b/scope.go
@@ -27,6 +27,8 @@ import (
"reflect"
"sort"
"time"
+
+ "go.uber.org/dig/internal/digclock"
)
// A ScopeOption modifies the default behavior of Scope; currently,
@@ -90,6 +92,9 @@ type Scope struct {
// All the child scopes of this Scope.
childScopes []*Scope
+
+ // clockSrc stores the source of time. Defaults to system clock.
+ clockSrc digclock.Clock
}
func newScope() *Scope {
@@ -102,6 +107,7 @@ func newScope() *Scope {
decoratedGroups: make(map[key]reflect.Value),
invokerFn: defaultInvoker,
rand: rand.New(rand.NewSource(time.Now().UnixNano())),
+ clockSrc: digclock.System,
}
s.gh = newGraphHolder(s)
return s
@@ -267,6 +273,10 @@ func (s *Scope) invoker() invokerFn {
return s.invokerFn
}
+func (s *Scope) clock() digclock.Clock {
+ return s.clockSrc
+}
+
// adds a new graphNode to this Scope and all of its descendent
// scope.
func (s *Scope) newGraphNode(wrapped interface{}, orders map[*Scope]int) {