pkg/leakybucket: replace global variables with injected collector

mmetc · mmetc · commit c5eb1c779c4f · 2026-01-14T13:46:08.000+01:00
diff --git a/cmd/crowdsec-cli/cliexplain/explain.go b/cmd/crowdsec-cli/cliexplain/explain.go
@@ -135,6 +135,7 @@ func (cli *cliExplain) run(ctx context.Context) error {
 
 	defer func() {
 		if cli.flags.noClean {
+			fmt.Fprintf(os.Stdout, "Not removing dump directory: %s\n", dir)
 			return
 		}
 
diff --git a/cmd/crowdsec/crowdsec.go b/cmd/crowdsec/crowdsec.go
@@ -19,6 +19,7 @@ import (
 	"github.com/crowdsecurity/crowdsec/pkg/csconfig"
 	"github.com/crowdsecurity/crowdsec/pkg/cwhub"
 	"github.com/crowdsecurity/crowdsec/pkg/exprhelpers"
+	"github.com/crowdsecurity/crowdsec/pkg/leakybucket"
 	"github.com/crowdsecurity/crowdsec/pkg/metrics"
 	"github.com/crowdsecurity/crowdsec/pkg/parser"
 	"github.com/crowdsecurity/crowdsec/pkg/pipeline"
@@ -81,12 +82,12 @@ func startParserRoutines(ctx context.Context, g *errgroup.Group, cConfig *csconf
 	}
 }
 
-func startBucketRoutines(ctx context.Context, g *errgroup.Group, cConfig *csconfig.Config) {
+func startBucketRoutines(ctx context.Context, g *errgroup.Group, cConfig *csconfig.Config, pourCollector *leakybucket.PourCollector) {
 	for idx := range cConfig.Crowdsec.BucketsRoutinesCount {
 		log.WithField("idx", idx).Info("Starting bucket routine")
 		g.Go(func() error {
 			defer trace.CatchPanic("crowdsec/runPour/"+strconv.Itoa(idx))
-			runPour(ctx, inEvents, holders, buckets, cConfig)
+			runPour(ctx, inEvents, holders, buckets, cConfig, pourCollector)
 			return nil
 		})
 	}
@@ -135,12 +136,12 @@ func startLPMetrics(ctx context.Context, cConfig *csconfig.Config, apiClient *ap
 }
 
 // runCrowdsec starts the log processor service
-func runCrowdsec(ctx context.Context, g *errgroup.Group, cConfig *csconfig.Config, parsers *parser.Parsers, hub *cwhub.Hub, datasources []acquisitionTypes.DataSource) error {
+func runCrowdsec(ctx context.Context, g *errgroup.Group, cConfig *csconfig.Config, parsers *parser.Parsers, hub *cwhub.Hub, datasources []acquisitionTypes.DataSource, pourCollector *leakybucket.PourCollector) error {
 	inEvents = make(chan pipeline.Event)
 	logLines = make(chan pipeline.Event)
 
 	startParserRoutines(ctx, g, cConfig, parsers)
-	startBucketRoutines(ctx, g, cConfig)
+	startBucketRoutines(ctx, g, cConfig, pourCollector)
 
 	apiClient, err := apiclient.GetLAPIClient()
 	if err != nil {
@@ -165,7 +166,7 @@ func runCrowdsec(ctx context.Context, g *errgroup.Group, cConfig *csconfig.Confi
 }
 
 // serveCrowdsec wraps the log processor service
-func serveCrowdsec(ctx context.Context, parsers *parser.Parsers, cConfig *csconfig.Config, hub *cwhub.Hub, datasources []acquisitionTypes.DataSource, agentReady chan bool) {
+func serveCrowdsec(ctx context.Context, parsers *parser.Parsers, cConfig *csconfig.Config, hub *cwhub.Hub, datasources []acquisitionTypes.DataSource, agentReady chan bool, pourCollector *leakybucket.PourCollector) {
 	cctx, cancel := context.WithCancel(ctx)
 
 	var g errgroup.Group
@@ -180,7 +181,7 @@ func serveCrowdsec(ctx context.Context, parsers *parser.Parsers, cConfig *csconf
 
 			agentReady <- true
 
-			if err := runCrowdsec(cctx, &g, cConfig, parsers, hub, datasources); err != nil {
+			if err := runCrowdsec(cctx, &g, cConfig, parsers, hub, datasources, pourCollector); err != nil {
 				log.Fatalf("unable to start crowdsec routines: %s", err)
 			}
 		}()
@@ -201,7 +202,7 @@ func serveCrowdsec(ctx context.Context, parsers *parser.Parsers, cConfig *csconf
 		if flags.DumpDir != "" {
 			log.Debugf("Dumping parser+bucket states to %s", flags.DumpDir)
 
-			if err := dumpAllStates(flags.DumpDir); err != nil {
+			if err := dumpAllStates(flags.DumpDir, pourCollector); err != nil {
 				log.Fatal(err)
 			}
 
diff --git a/cmd/crowdsec/dump.go b/cmd/crowdsec/dump.go
@@ -7,11 +7,12 @@ import (
 
 	"gopkg.in/yaml.v3"
 
-	leaky "github.com/crowdsecurity/crowdsec/pkg/leakybucket"
+	"github.com/crowdsecurity/crowdsec/pkg/leakybucket"
 	"github.com/crowdsecurity/crowdsec/pkg/parser"
+	"github.com/crowdsecurity/crowdsec/pkg/pipeline"
 )
 
-func dumpAllStates(dir string) error {
+func dumpAllStates(dir string, pourCollector *leakybucket.PourCollector) error {
 	err := os.MkdirAll(dir, 0o755)
 	if err != nil {
 		return err
@@ -25,13 +26,30 @@ func dumpAllStates(dir string) error {
 		return fmt.Errorf("dumping bucket overflow state: %w", err)
 	}
 
-	if err := dumpState(dir, "bucketpour-dump.yaml", leaky.BucketPourCache); err != nil {
+	if err := dumpCollector(dir, "bucketpour-dump.yaml", pourCollector); err != nil {
 		return fmt.Errorf("dumping bucket pour state: %w", err)
 	}
 
 	return nil
 }
 
+type Collector interface {
+	Snapshot() map[string][]pipeline.Event
+}
+
+func dumpCollector(dir, name string, collector Collector) error {
+	if collector == nil {
+		return nil
+	}
+
+	out, err := yaml.Marshal(collector.Snapshot())
+	if err != nil {
+		return err
+	}
+
+	return os.WriteFile(filepath.Join(dir, name), out, 0o644)
+}
+
 func dumpState(dir, name string, obj any) error {
 	out, err := yaml.Marshal(obj)
 	if err != nil {
diff --git a/cmd/crowdsec/main.go b/cmd/crowdsec/main.go
@@ -213,7 +213,13 @@ func run(flags Flags) error {
 		return err
 	}
 
-	return StartRunSvc(ctx, cConfig)
+	var pourCollector *leakybucket.PourCollector
+
+	if flags.DumpDir != "" {
+		pourCollector = leakybucket.NewPourCollector()
+	}
+
+	return StartRunSvc(ctx, cConfig, pourCollector)
 }
 
 func main() {
diff --git a/cmd/crowdsec/pour.go b/cmd/crowdsec/pour.go
@@ -35,7 +35,7 @@ func triggerGC(parsed pipeline.Event, buckets *leaky.Buckets, cConfig *csconfig.
 	leaky.GarbageCollectBuckets(*z, buckets)
 }
 
-func runPour(ctx context.Context, input chan pipeline.Event, holders []leaky.BucketFactory, buckets *leaky.Buckets, cConfig *csconfig.Config) {
+func runPour(ctx context.Context, input chan pipeline.Event, holders []leaky.BucketFactory, buckets *leaky.Buckets, cConfig *csconfig.Config, pourCollector *leaky.PourCollector) {
 	count := 0
 
 	for {
@@ -52,8 +52,7 @@ func runPour(ctx context.Context, input chan pipeline.Event, holders []leaky.Buc
 				triggerGC(parsed, buckets, cConfig)
 			}
 			// here we can bucketify with parsed
-			track := flags.DumpDir != ""
-			poured, err := leaky.PourItemToHolders(ctx, parsed, holders, buckets, track)
+			poured, err := leaky.PourItemToHolders(ctx, parsed, holders, buckets, pourCollector)
 			if err != nil {
 				log.Warningf("bucketify failed for: %v with %s", parsed, err)
 				continue
diff --git a/cmd/crowdsec/run_in_svc.go b/cmd/crowdsec/run_in_svc.go
@@ -16,13 +16,14 @@ import (
 	"github.com/crowdsecurity/crowdsec/pkg/csconfig"
 	"github.com/crowdsecurity/crowdsec/pkg/database"
 	"github.com/crowdsecurity/crowdsec/pkg/fflag"
+	"github.com/crowdsecurity/crowdsec/pkg/leakybucket"
 )
 
 func isWindowsService() (bool, error) {
 	return false, nil
 }
 
-func StartRunSvc(ctx context.Context, cConfig *csconfig.Config) error {
+func StartRunSvc(ctx context.Context, cConfig *csconfig.Config, pourCollector *leakybucket.PourCollector) error {
 	defer trace.CatchPanic("crowdsec/StartRunSvc")
 
 	// Always try to stop CPU profiling to avoid passing flags around
@@ -63,5 +64,5 @@ func StartRunSvc(ctx context.Context, cConfig *csconfig.Config) error {
 		}()
 	}
 
-	return Serve(ctx, cConfig, agentReady)
+	return Serve(ctx, cConfig, agentReady, pourCollector)
 }
diff --git a/cmd/crowdsec/serve.go b/cmd/crowdsec/serve.go
@@ -22,6 +22,7 @@ import (
 	"github.com/crowdsecurity/crowdsec/pkg/cwhub"
 	"github.com/crowdsecurity/crowdsec/pkg/database"
 	"github.com/crowdsecurity/crowdsec/pkg/exprhelpers"
+	"github.com/crowdsecurity/crowdsec/pkg/leakybucket"
 	"github.com/crowdsecurity/crowdsec/pkg/pipeline"
 )
 
@@ -33,6 +34,12 @@ func reloadHandler(ctx context.Context, _ os.Signal) (*csconfig.Config, error) {
 	crowdsecTomb = tomb.Tomb{}
 	pluginTomb = tomb.Tomb{}
 
+	var pourCollector *leakybucket.PourCollector
+
+	if flags.DumpDir != "" {
+		pourCollector = leakybucket.NewPourCollector()
+	}
+
 	cConfig, err := LoadConfig(flags.ConfigFile, flags.DisableAgent, flags.DisableAPI, false)
 	if err != nil {
 		return nil, err
@@ -77,7 +84,7 @@ func reloadHandler(ctx context.Context, _ os.Signal) (*csconfig.Config, error) {
 		}
 
 		agentReady := make(chan bool, 1)
-		serveCrowdsec(ctx, csParsers, cConfig, hub, datasources, agentReady)
+		serveCrowdsec(ctx, csParsers, cConfig, hub, datasources, agentReady, pourCollector)
 	}
 
 	log.Info("Reload is finished")
@@ -301,7 +308,7 @@ func HandleSignals(ctx context.Context, cConfig *csconfig.Config) error {
 	return err
 }
 
-func Serve(ctx context.Context, cConfig *csconfig.Config, agentReady chan bool) error {
+func Serve(ctx context.Context, cConfig *csconfig.Config, agentReady chan bool, pourCollector *leakybucket.PourCollector) error {
 	acquisTomb = tomb.Tomb{}
 	outputsTomb = tomb.Tomb{}
 	apiTomb = tomb.Tomb{}
@@ -374,7 +381,7 @@ func Serve(ctx context.Context, cConfig *csconfig.Config, agentReady chan bool)
 
 		// if it's just linting, we're done
 		if !flags.TestMode {
-			serveCrowdsec(ctx, csParsers, cConfig, hub, datasources, agentReady)
+			serveCrowdsec(ctx, csParsers, cConfig, hub, datasources, agentReady, pourCollector)
 		} else {
 			agentReady <- true
 		}
diff --git a/pkg/leakybucket/buckets_test.go b/pkg/leakybucket/buckets_test.go
@@ -184,7 +184,7 @@ func testFile(t *testing.T, file string, holders []BucketFactory, response chan
 		in.ExpectMode = pipeline.TIMEMACHINE
 		log.Infof("Buckets input : %s", spew.Sdump(in))
 
-		ok, err := PourItemToHolders(ctx, in, holders, buckets, false)
+		ok, err := PourItemToHolders(ctx, in, holders, buckets, nil)
 		if err != nil {
 			t.Fatalf("Failed to pour : %s", err)
 		}
@@ -207,7 +207,7 @@ POLL_AGAIN:
 			results = append(results, ret)
 			if ret.Overflow.Reprocess {
 				log.Errorf("Overflow being reprocessed.")
-				ok, err := PourItemToHolders(ctx, ret, holders, buckets, false)
+				ok, err := PourItemToHolders(ctx, ret, holders, buckets, nil)
 				if err != nil {
 					t.Fatalf("Failed to pour : %s", err)
 				}
diff --git a/pkg/leakybucket/collector.go b/pkg/leakybucket/collector.go
@@ -0,0 +1,45 @@
+package leakybucket
+
+import (
+	"sync"
+	
+	"github.com/crowdsecurity/crowdsec/pkg/pipeline"
+)
+
+type PourCollector struct {
+	mu sync.Mutex
+	m  map[string][]pipeline.Event
+}
+
+func NewPourCollector() *PourCollector {
+	return &PourCollector{
+		m: make(map[string][]pipeline.Event),
+	}
+}
+
+func (c *PourCollector) Add(key string, evt pipeline.Event) {
+	if c == nil {
+		return
+	}
+	c.mu.Lock()
+	c.m[key] = append(c.m[key], evt)
+	c.mu.Unlock()
+}
+
+// Snapshot returns a shallow copy of the map and slices.
+// The caller must not mutate the events.
+func (c *PourCollector) Snapshot() map[string][]pipeline.Event {
+	if c == nil {
+		return nil
+	}
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	out := make(map[string][]pipeline.Event, len(c.m))
+	for k, v := range c.m {
+		tmp := make([]pipeline.Event, len(v))
+		copy(tmp, v)
+		out[k] = tmp
+	}
+	return out
+}
diff --git a/pkg/leakybucket/manager_run.go b/pkg/leakybucket/manager_run.go
@@ -19,8 +19,6 @@ import (
 
 var (
 	serialized      map[string]Leaky
-	BucketPourCache map[string][]pipeline.Event = make(map[string][]pipeline.Event)
-	bucketPourMu    sync.Mutex
 )
 
 /*
@@ -74,7 +72,7 @@ func GarbageCollectBuckets(deadline time.Time, buckets *Buckets) {
 	}
 }
 
-func PourItemToBucket(ctx context.Context, bucket *Leaky, holder BucketFactory, buckets *Buckets, parsed *pipeline.Event, track bool) (bool, error) {
+func PourItemToBucket(ctx context.Context, bucket *Leaky, holder BucketFactory, buckets *Buckets, parsed *pipeline.Event, collector *PourCollector) (bool, error) {
 	var sent bool
 	var buckey = bucket.Mapkey
 	var err error
@@ -142,12 +140,9 @@ func PourItemToBucket(ctx context.Context, bucket *Leaky, holder BucketFactory,
 		select {
 		case bucket.In <- parsed:
 			// holder.logger.Tracef("Successfully sent !")
-			if track {
+			if collector != nil {
 				evt := deepcopy.Copy(*parsed).(pipeline.Event)
-
-				bucketPourMu.Lock()
-				BucketPourCache[bucket.Name] = append(BucketPourCache[bucket.Name], evt)
-				bucketPourMu.Unlock()
+				collector.Add(bucket.Name, evt)
 			}
 			sent = true
 			continue
@@ -203,15 +198,12 @@ func LoadOrStoreBucketFromHolder(ctx context.Context, partitionKey string, bucke
 
 var orderEvent map[string]*sync.WaitGroup
 
-func PourItemToHolders(ctx context.Context, parsed pipeline.Event, holders []BucketFactory, buckets *Buckets, track bool) (bool, error) {
+func PourItemToHolders(ctx context.Context, parsed pipeline.Event, holders []BucketFactory, buckets *Buckets, collector *PourCollector) (bool, error) {
 	var ok, condition, poured bool
 
-	if track {
+	if collector != nil {
 		evt := deepcopy.Copy(parsed).(pipeline.Event)
-
-		bucketPourMu.Lock()
-		BucketPourCache["OK"] = append(BucketPourCache["OK"], evt)
-		bucketPourMu.Unlock()
+		collector.Add("OK", evt)
 	}
 	// find the relevant holders (scenarios)
 	for idx := range holders {
@@ -274,7 +266,7 @@ func PourItemToHolders(ctx context.Context, parsed pipeline.Event, holders []Buc
 			orderEvent[buckey].Add(1)
 		}
 
-		ok, err := PourItemToBucket(ctx, bucket, holders[idx], buckets, &parsed, track)
+		ok, err := PourItemToBucket(ctx, bucket, holders[idx], buckets, &parsed, collector)
 
 		if bucket.orderEvent {
 			orderEvent[buckey].Wait()
diff --git a/pkg/leakybucket/manager_run_test.go b/pkg/leakybucket/manager_run_test.go
@@ -86,7 +86,7 @@ func TestGCandDump(t *testing.T) {
 
 	in := pipeline.Event{Parsed: map[string]string{"something": "something"}}
 	// pour an item that will go to leaky + counter
-	ok, err := PourItemToHolders(ctx, in, Holders, buckets, false)
+	ok, err := PourItemToHolders(ctx, in, Holders, buckets, nil)
 	if err != nil {
 		t.Fatalf("while pouring item: %s", err)
 	}
@@ -157,7 +157,7 @@ func TestShutdownBuckets(t *testing.T) {
 	in := pipeline.Event{Parsed: map[string]string{"something": "something"}}
 	// pour an item that will go to leaky + counter
 	ctx, cancel := context.WithCancel(t.Context())
-	ok, err := PourItemToHolders(ctx, in, Holders, buckets, false)
+	ok, err := PourItemToHolders(ctx, in, Holders, buckets, nil)
 	if err != nil {
 		t.Fatalf("while pouring item : %s", err)
 	}

Original file line number	Diff line number	Diff line change
`@@ -135,6 +135,7 @@ func (cli *cliExplain) run(ctx context.Context) error {`
`135`	`135`
`136`	`136`	`defer func() {`
`137`	`137`	`if cli.flags.noClean {`
	`138`	`+ fmt.Fprintf(os.Stdout, "Not removing dump directory: %s\n", dir)`
`138`	`139`	`return`
`139`	`140`	`}`
`140`	`141`
Original file line number	Diff line number	Diff line change
`@@ -213,7 +213,13 @@ func run(flags Flags) error {`
`213`	`213`	`return err`
`214`	`214`	`}`
`215`	`215`
`216`		`- return StartRunSvc(ctx, cConfig)`
	`216`	`+ var pourCollector *leakybucket.PourCollector`
	`217`	`+`
	`218`	`+ if flags.DumpDir != "" {`
	`219`	`+ pourCollector = leakybucket.NewPourCollector()`
	`220`	`+ }`
	`221`	`+`
	`222`	`+ return StartRunSvc(ctx, cConfig, pourCollector)`
`217`	`223`	`}`
`218`	`224`
`219`	`225`	`func main() {`
Original file line number	Diff line number	Diff line change
`@@ -16,13 +16,14 @@ import (`
`16`	`16`	`"github.com/crowdsecurity/crowdsec/pkg/csconfig"`
`17`	`17`	`"github.com/crowdsecurity/crowdsec/pkg/database"`
`18`	`18`	`"github.com/crowdsecurity/crowdsec/pkg/fflag"`
	`19`	`+ "github.com/crowdsecurity/crowdsec/pkg/leakybucket"`
`19`	`20`	`)`
`20`	`21`
`21`	`22`	`func isWindowsService() (bool, error) {`
`22`	`23`	`return false, nil`
`23`	`24`	`}`
`24`	`25`
`25`		`-func StartRunSvc(ctx context.Context, cConfig *csconfig.Config) error {`
	`26`	`+func StartRunSvc(ctx context.Context, cConfig csconfig.Config, pourCollector leakybucket.PourCollector) error {`
`26`	`27`	`defer trace.CatchPanic("crowdsec/StartRunSvc")`
`27`	`28`
`28`	`29`	`// Always try to stop CPU profiling to avoid passing flags around`
`@@ -63,5 +64,5 @@ func StartRunSvc(ctx context.Context, cConfig *csconfig.Config) error {`
`63`	`64`	`}()`
`64`	`65`	`}`
`65`	`66`
`66`		`- return Serve(ctx, cConfig, agentReady)`
	`67`	`+ return Serve(ctx, cConfig, agentReady, pourCollector)`
`67`	`68`	`}`