add counter reset support

pkg/genlib/config/config.go

@@ -39,15 +39,55 @@ type Config struct {
 }
 
 type ConfigField struct {
-	Name        string        `config:"name"`
-	Fuzziness   float64       `config:"fuzziness"`
-	Range       Range         `config:"range"`
-	Cardinality int           `config:"cardinality"`
-	Period      time.Duration `config:"period"`
-	Enum        []string      `config:"enum"`
-	ObjectKeys  []string      `config:"object_keys"`
-	Value       any           `config:"value"`
-	Counter     bool          `config:"counter"`
+	Name         string        `config:"name"`
+	Fuzziness    float64       `config:"fuzziness"`
+	Range        Range         `config:"range"`
+	Cardinality  int           `config:"cardinality"`
+	Period       time.Duration `config:"period"`
+	Enum         []string      `config:"enum"`
+	ObjectKeys   []string      `config:"object_keys"`
+	Value        any           `config:"value"`
+	Counter      bool          `config:"counter"`
+	CounterReset *CounterReset `config:"counter_reset"`
+}
+
+const (
+	CounterResetStrategyRandom        string = "random"
+	CounterResetStrategyProbabilistic string = "probabilistic"
+	CounterResetStrategyAfterN        string = "after_n"
+)
+
+type CounterReset struct {
+	Strategy    string  `config:"strategy"`
+	Probability *uint64 `config:"probability"`
+	ResetAfterN *uint64 `config:"reset_after_n"`
+}
+
+func (cf ConfigField) ValidateCounterResetStrategy() error {
+	if cf.Counter && cf.CounterReset != nil &&
+		cf.CounterReset.Strategy != CounterResetStrategyRandom &&
+		cf.CounterReset.Strategy != CounterResetStrategyProbabilistic &&
+		cf.CounterReset.Strategy != CounterResetStrategyAfterN {
+		return errors.New("counter_reset strategy must be one of 'random', 'probabilistic', 'after_n'")
+	}
+
+	return nil
+}
+
+func (cf ConfigField) ValidateCounterResetAfterN() error {
+	if cf.Counter && cf.CounterReset != nil && cf.CounterReset.Strategy == CounterResetStrategyAfterN && cf.CounterReset.ResetAfterN == nil {
+		return errors.New("counter_reset after_n requires 'reset_after_n' value to be set")
+	}
+
+	return nil
+}
+
+func (cf ConfigField) ValidateCounterResetProbabilistic() error {
+	if cf.Counter && cf.CounterReset != nil && cf.CounterReset.Strategy == CounterResetStrategyProbabilistic && cf.CounterReset.Probability == nil {
+		return errors.New("counter_reset probabilistic requires 'probability' value to be set")
+	}
+
+	return nil
 }
 
 func (cf ConfigField) ValidForDateField() error {

pkg/genlib/generator_interface.go

@@ -82,7 +82,8 @@ type genState struct {
 	// previous cardinality value cache; necessary for cardinality
 	prevCacheCardinality map[string][]any
 	// internal buffer pool to decrease load on GC
-	pool sync.Pool
+	pool         sync.Pool
+	counterReset bool
 }
 
 func newGenState() *genState {
@@ -99,7 +100,6 @@ func newGenState() *genState {
 }
 
 func bindField(cfg Config, field Field, fieldMap map[string]any, withReturn bool) error {
-
 	// Check for hardcoded field value
 	if len(field.Value) > 0 {
 		if withReturn {
@@ -194,7 +194,6 @@ func bindByType(cfg Config, field Field, fieldMap map[string]any) (err error) {
 }
 
 func bindByTypeWithReturn(cfg Config, field Field, fieldMap map[string]any) (err error) {
-
 	fieldCfg, _ := cfg.GetField(field.Name)
 
 	switch field.Type {
@@ -971,6 +970,18 @@ func bindLongWithReturn(fieldCfg ConfigField, field Field, fieldMap map[string]a
 		return err
 	}
 
+	if err := fieldCfg.ValidateCounterResetStrategy(); err != nil {
+		return err
+	}
+
+	if err := fieldCfg.ValidateCounterResetAfterN(); err != nil {
+		return err
+	}
+
+	if err := fieldCfg.ValidateCounterResetProbabilistic(); err != nil {
+		return err
+	}
+
 	if len(fieldCfg.Enum) > 0 {
 		var emitF emitF
 		idx := customRand.Intn(len(fieldCfg.Enum))
@@ -1008,6 +1019,27 @@ func bindLongWithReturn(fieldCfg ConfigField, field Field, fieldMap map[string]a
 				dummyInt = fuzzyIntCounter(previous, fieldCfg.Fuzziness)
 			}
 
+			if fieldCfg.CounterReset != nil {
+				switch fieldCfg.CounterReset.Strategy {
+				case config.CounterResetStrategyRandom:
+					// 50% chance to reset
+					if customRand.Intn(2) == 0 {
+						dummyInt = 0
+					}
+				case config.CounterResetStrategyProbabilistic:
+					// Probability% chance to reset
+					if customRand.Intn(100) < int(*fieldCfg.CounterReset.Probability) {
+						dummyInt = 0
+					}
+				case config.CounterResetStrategyAfterN:
+					// Reset after N
+					if !state.counterReset && state.counter >= *fieldCfg.CounterReset.ResetAfterN {
+						dummyInt = 0
+						state.counterReset = true
+					}
+				}
+			}
+
 			state.prevCache[field.Name] = dummyInt
 			return dummyInt
 		}

pkg/genlib/generator_with_text_template_test.go

@@ -755,6 +755,64 @@ func Test_FieldIPWithTextTemplate(t *testing.T) {
 	}
 }
 
+func Test_FieldLongCounterResetAfterN5WithTextTemplate(t *testing.T) {
+	fld := Field{
+		Name: "counter_reset_test",
+		Type: FieldTypeLong,
+	}
+
+	afterN := 5
+
+	template := []byte(`{{$counter_reset_test := generate "counter_reset_test"}}{"counter_reset_test":"{{$counter_reset_test}}"}`)
+	configYaml := []byte(fmt.Sprintf(`fields:
+- name: counter_reset_test
+  counter: true
+  counter_reset:
+    strategy: after_n
+    reset_after_n: %d`, afterN))
+	t.Logf("with template: %s", string(template))
+
+	cfg, err := config.LoadConfigFromYaml(configYaml)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	g := makeGeneratorWithTextTemplate(t, cfg, []Field{fld}, template, 10)
+
+	var buf bytes.Buffer
+
+	nSpins := int64(10)
+
+	var shouldReset bool
+
+	for i := int64(0); i < nSpins; i++ {
+		if err := g.Emit(&buf); err != nil {
+			t.Fatal(err)
+		}
+
+		m := unmarshalJSONT[string](t, buf.Bytes())
+		buf.Reset()
+
+		if len(m) != 1 {
+			t.Errorf("Expected map size 1, got %d", len(m))
+		}
+
+		v, ok := m[fld.Name]
+		if !ok {
+			t.Errorf("Missing key %v", fld.Name)
+		}
+
+		if i >= int64(afterN) && !shouldReset {
+			if v != "0" {
+				t.Errorf("Expected counter to reset to 0, got %v", v)
+			}
+			shouldReset = true
+		}
+
+		t.Logf("counter value: %v", v)
+	}
+}
+
 func Test_FieldFloatsWithTextTemplate(t *testing.T) {
 	_testNumericWithTextTemplate[float64](t, FieldTypeDouble)
 	_testNumericWithTextTemplate[float32](t, FieldTypeFloat)