From 556fd59b42f68a2fb1f84957741b72811c714e51 Mon Sep 17 00:00:00 2001
From: Mikhail Mazurskiy <126021+ash2k@users.noreply.github.com>
Date: Thu, 2 Nov 2023 00:45:39 +1100
Subject: [PATCH] time: per-thread time.now() function (#517)

---
 lib/time/time.go      | 35 ++++++++++++++++--
 lib/time/time_test.go | 82 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 115 insertions(+), 2 deletions(-)
 create mode 100644 lib/time/time_test.go

diff --git a/lib/time/time.go b/lib/time/time.go
index 0f781420..530716cc 100644
--- a/lib/time/time.go
+++ b/lib/time/time.go
@@ -6,6 +6,7 @@
 package time // import "go.starlark.net/lib/time"
 
 import (
+	"errors"
 	"fmt"
 	"sort"
 	"time"
@@ -68,11 +69,29 @@ var Module = &starlarkstruct.Module{
 	},
 }
 
-// NowFunc is a function that generates the current time. Intentionally exported
+// NowFunc is a function that reports the current time. Intentionally exported
 // so that it can be overridden, for example by applications that require their
 // Starlark scripts to be fully deterministic.
+//
+// Deprecated: avoid updating this global variable
+// and instead use SetNow on each thread to set its clock function.
 var NowFunc = time.Now
 
+const contextKey = "time.now"
+
+// SetNow sets the thread's optional clock function.
+// If non-nil, it will be used in preference to NowFunc when the
+// thread requests the current time by executing a call to time.now.
+func SetNow(thread *starlark.Thread, nowFunc func() (time.Time, error)) {
+	thread.SetLocal(contextKey, nowFunc)
+}
+
+// Now returns the clock function previously associated with this thread.
+func Now(thread *starlark.Thread) func() (time.Time, error) {
+	nowFunc, _ := thread.Local(contextKey).(func() (time.Time, error))
+	return nowFunc
+}
+
 func parseDuration(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
 	var d Duration
 	err := starlark.UnpackPositionalArgs("parse_duration", args, kwargs, 1, &d)
@@ -129,7 +148,19 @@ func fromTimestamp(thread *starlark.Thread, _ *starlark.Builtin, args starlark.T
 }
 
 func now(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
-	return Time(NowFunc()), nil
+	nowErrFunc := Now(thread)
+	if nowErrFunc != nil {
+		t, err := nowErrFunc()
+		if err != nil {
+			return nil, err
+		}
+		return Time(t), nil
+	}
+	nowFunc := NowFunc
+	if nowFunc == nil {
+		return nil, errors.New("time.now() is not available")
+	}
+	return Time(nowFunc()), nil
 }
 
 // Duration is a Starlark representation of a duration.
diff --git a/lib/time/time_test.go b/lib/time/time_test.go
new file mode 100644
index 00000000..b799953b
--- /dev/null
+++ b/lib/time/time_test.go
@@ -0,0 +1,82 @@
+package time
+
+import (
+	"errors"
+	"testing"
+	"time"
+
+	"go.starlark.net/starlark"
+)
+
+func TestPerThreadNowReturnsCorrectTime(t *testing.T) {
+	th := &starlark.Thread{}
+	date := time.Date(1, 2, 3, 4, 5, 6, 7, time.UTC)
+	SetNow(th, func() (time.Time, error) {
+		return date, nil
+	})
+
+	res, err := starlark.Call(th, Module.Members["now"], nil, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	retTime := time.Time(res.(Time))
+
+	if !retTime.Equal(date) {
+		t.Fatal("Expected time to be equal", retTime, date)
+	}
+}
+
+func TestPerThreadNowReturnsError(t *testing.T) {
+	th := &starlark.Thread{}
+	e := errors.New("no time")
+	SetNow(th, func() (time.Time, error) {
+		return time.Time{}, e
+	})
+
+	_, err := starlark.Call(th, Module.Members["now"], nil, nil)
+	if !errors.Is(err, e) {
+		t.Fatal("Expected equal error", e, err)
+	}
+}
+
+func TestGlobalNowReturnsCorrectTime(t *testing.T) {
+	th := &starlark.Thread{}
+
+	oldNow := NowFunc
+	defer func() {
+		NowFunc = oldNow
+	}()
+
+	date := time.Date(1, 2, 3, 4, 5, 6, 7, time.UTC)
+	NowFunc = func() time.Time {
+		return date
+	}
+
+	res, err := starlark.Call(th, Module.Members["now"], nil, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	retTime := time.Time(res.(Time))
+
+	if !retTime.Equal(date) {
+		t.Fatal("Expected time to be equal", retTime, date)
+	}
+}
+
+func TestGlobalNowReturnsErrorWhenNil(t *testing.T) {
+	th := &starlark.Thread{}
+
+	oldNow := NowFunc
+	defer func() {
+		NowFunc = oldNow
+	}()
+
+	NowFunc = nil
+
+	_, err := starlark.Call(th, Module.Members["now"], nil, nil)
+	if err == nil {
+		t.Fatal("Expected to get an error")
+	}
+}