Expose time globals to Lua raters and generators

generator/lua.go

@@ -451,6 +451,12 @@ func (lg *luagen) Gen(item *config.GenQueueItem) error {
 				L.SetGlobal("earliest", lua.LNumber(float64(item.Earliest.UnixNano())/float64(time.Second)))
 				L.SetGlobal("latest", lua.LNumber(float64(item.Latest.UnixNano())/float64(time.Second)))
 				L.SetGlobal("now", lua.LNumber(float64(item.Now.UnixNano())/float64(time.Second)))
+				if !s.BeginParsed.IsZero() {
+					L.SetGlobal("beginTime", lua.LNumber(float64(s.BeginParsed.UnixNano())/float64(time.Second)))
+				}
+				if !s.EndParsed.IsZero() {
+					L.SetGlobal("endTime", lua.LNumber(float64(s.EndParsed.UnixNano())/float64(time.Second)))
+				}
 
 				// Register functions
 				L.SetGlobal("sleep", L.NewFunction(sleep))
@@ -484,6 +490,12 @@ func (lg *luagen) Gen(item *config.GenQueueItem) error {
 	L.SetGlobal("earliest", lua.LNumber(float64(item.Earliest.UnixNano())/float64(time.Second)))
 	L.SetGlobal("latest", lua.LNumber(float64(item.Latest.UnixNano())/float64(time.Second)))
 	L.SetGlobal("now", lua.LNumber(float64(item.Now.UnixNano())/float64(time.Second)))
+	if !s.BeginParsed.IsZero() {
+		L.SetGlobal("beginTime", lua.LNumber(float64(s.BeginParsed.UnixNano())/float64(time.Second)))
+	}
+	if !s.EndParsed.IsZero() {
+		L.SetGlobal("endTime", lua.LNumber(float64(s.EndParsed.UnixNano())/float64(time.Second)))
+	}
 
 	// log.Debugf("Calling DoString for %# v", s.CustomGenerator.Script)
 	var f *lua.LFunction

generator/lua_test.go

@@ -1,6 +1,7 @@
 package generator
 
 import (
+	"fmt"
 	"math/rand"
 	"os"
 	"path/filepath"
@@ -419,3 +420,27 @@ func runLuaGen(t *testing.T, s *config.Sample, gen *luagen) (chan *config.OutQue
 	}()
 	return oq, err
 }
+
+func TestBeginEndTimeExposed(t *testing.T) {
+	config.ResetConfig()
+
+	os.Setenv("GOGEN_HOME", "..")
+	os.Setenv("GOGEN_ALWAYS_REFRESH", "")
+	home := ".."
+	os.Setenv("GOGEN_FULLCONFIG", filepath.Join(home, "tests", "generator", "luaapi_time.yml"))
+
+	c := config.NewConfig()
+	s := c.FindSampleByName("beginEndTime")
+
+	// Set BeginParsed and EndParsed on the sample
+	loc, _ := time.LoadLocation("Local")
+	s.BeginParsed = time.Date(2001, 10, 20, 11, 0, 0, 0, loc)
+	s.EndParsed = time.Date(2001, 10, 20, 13, 0, 0, 0, loc)
+
+	beginEpoch := s.BeginParsed.Unix()
+	endEpoch := s.EndParsed.Unix()
+	expected := fmt.Sprintf("%d-%d", beginEpoch, endEpoch)
+
+	gen := new(luagen)
+	testLuaGen(t, s, gen, expected)
+}

rater/script.go

@@ -17,8 +17,10 @@ type ScriptRater struct {
 	luaState *lua.LTable
 }
 
-// GetRate acts as a general method for EventRate and TokenRate
-func (sr *ScriptRater) getRate(now time.Time) float64 {
+// getRate acts as a general method for EventRate and TokenRate.
+// It exposes now, beginTime, and endTime as Lua globals so scripts
+// can make time-aware decisions (e.g. operating hours, week-based ramp-up).
+func (sr *ScriptRater) getRate(s *config.Sample, now time.Time) float64 {
 	if sr.luaState == nil {
 		sr.luaState = new(lua.LTable)
 		for k, v := range sr.c.Init {
@@ -29,6 +31,15 @@ func (sr *ScriptRater) getRate(now time.Time) float64 {
 	defer L.Close()
 	L.SetGlobal("state", sr.luaState)
 	L.SetGlobal("options", luar.New(L, sr.c.Options))
+	L.SetGlobal("now", lua.LNumber(float64(now.UnixNano())/float64(time.Second)))
+	if s != nil {
+		if !s.BeginParsed.IsZero() {
+			L.SetGlobal("beginTime", lua.LNumber(float64(s.BeginParsed.UnixNano())/float64(time.Second)))
+		}
+		if !s.EndParsed.IsZero() {
+			L.SetGlobal("endTime", lua.LNumber(float64(s.EndParsed.UnixNano())/float64(time.Second)))
+		}
+	}
 	if err := L.DoString(sr.c.Script); err != nil {
 		log.Errorf("Error executing script for rater '%s': %s", sr.c.Name, err)
 	}
@@ -37,10 +48,10 @@ func (sr *ScriptRater) getRate(now time.Time) float64 {
 
 // EventRate takes a given sample and current count and returns the rated count
 func (sr *ScriptRater) EventRate(s *config.Sample, now time.Time, count int) float64 {
-	return sr.getRate(now)
+	return sr.getRate(s, now)
 }
 
 // TokenRate takes a token and returns the rated value
 func (sr *ScriptRater) TokenRate(t config.Token, now time.Time) float64 {
-	return sr.getRate(now)
+	return sr.getRate(t.Parent, now)
 }

rater/script_test.go

@@ -28,3 +28,43 @@ func TestScriptRaterEventRate(t *testing.T) {
 	assert.True(t, assert.ObjectsAreEqual(r, s.Rater.(*ScriptRater).c))
 	assert.Equal(t, 2, ret)
 }
+
+func TestScriptRaterNowExposed(t *testing.T) {
+	config.ResetConfig()
+	os.Setenv("GOGEN_HOME", "..")
+	os.Setenv("GOGEN_ALWAYS_REFRESH", "1")
+	home := ".."
+	os.Setenv("GOGEN_FULLCONFIG", filepath.Join(home, "tests", "rater", "luarater_time.yml"))
+
+	c := config.NewConfig()
+	s := c.FindSampleByName("time_aware")
+	assert.Equal(t, "time_rater", s.RaterString)
+	// The script returns 3.0 if now > 0, else 1.0
+	ret := EventRate(s, time.Now(), 1)
+	assert.Equal(t, 3, ret)
+}
+
+func TestScriptRaterBeginTimeExposed(t *testing.T) {
+	config.ResetConfig()
+	os.Setenv("GOGEN_HOME", "..")
+	os.Setenv("GOGEN_ALWAYS_REFRESH", "1")
+	home := ".."
+	os.Setenv("GOGEN_FULLCONFIG", filepath.Join(home, "tests", "rater", "luarater_time.yml"))
+
+	c := config.NewConfig()
+	s := c.FindSampleByName("time_aware")
+	// Set BeginParsed so beginTime is exposed
+	s.BeginParsed = time.Date(2001, 10, 20, 12, 0, 0, 0, time.UTC)
+
+	// Override rater with one that checks beginTime
+	r := &ScriptRater{
+		c: &config.RaterConfig{
+			Name:   "begin_check",
+			Type:   "script",
+			Script: "if beginTime ~= nil and beginTime > 0 then return 5.0 else return 1.0 end",
+		},
+	}
+	s.Rater = r
+	ret := r.EventRate(s, time.Now(), 1)
+	assert.Equal(t, float64(5.0), ret)
+}

tests/generator/luaapi_time.yml

@@ -0,0 +1,15 @@
+generators:
+  - name: beginEndTime
+    script: |
+        line = getLine(0)
+        setToken("bt", tostring(math.floor(beginTime)))
+        setToken("et", tostring(math.floor(endTime)))
+        line = replaceTokens(line)
+        sendEvent(line)
+samples:
+  - name: beginEndTime
+    generator: beginEndTime
+    interval: 1
+    endIntervals: 1
+    lines:
+    - _raw: "$bt$-$et$"

tests/rater/luarater_time.yml

@@ -0,0 +1,15 @@
+global:
+  output:
+    outputter: buf
+samples:
+  - name: time_aware
+    rater: time_rater
+    count: 1
+    endIntervals: 1
+    lines:
+    - "_raw": test
+raters:
+  - name: time_rater
+    type: script
+    script: >
+        if now > 0 then return 3.0 else return 1.0 end