diff --git a/generator/lua.go b/generator/lua.go index ad94156..4b113fb 100644 --- a/generator/lua.go +++ b/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 diff --git a/generator/lua_test.go b/generator/lua_test.go index cc3534c..0771dea 100644 --- a/generator/lua_test.go +++ b/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) +} diff --git a/rater/script.go b/rater/script.go index dff22f1..54e3f7a 100644 --- a/rater/script.go +++ b/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) } diff --git a/rater/script_test.go b/rater/script_test.go index 307f35c..62ea431 100644 --- a/rater/script_test.go +++ b/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) +} diff --git a/tests/generator/luaapi_time.yml b/tests/generator/luaapi_time.yml new file mode 100644 index 0000000..f07cb35 --- /dev/null +++ b/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$" diff --git a/tests/rater/luarater_time.yml b/tests/rater/luarater_time.yml new file mode 100644 index 0000000..df41919 --- /dev/null +++ b/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