Add dynamic enable/disable of flow emission (#9995)

Author: caseydavenport

goldmane/pkg/aggregator/aggregator.go

@@ -87,7 +87,10 @@ type LogAggregator struct {
 	streamRequests chan streamRequest
 
 	// sink is a sink to send aggregated flows to.
-	sink Sink
+	sink bucketing.Sink
+
+	// sinkChan allows setting the sink asynchronously.
+	sinkChan chan bucketing.Sink
 
 	// recvChan is the channel to receive flow updates on.
 	recvChan chan *proto.FlowUpdate
@@ -124,6 +127,7 @@ func NewLogAggregator(opts ...Option) *LogAggregator {
 		listRequests:       make(chan listRequest),
 		streamRequests:     make(chan streamRequest),
 		recvChan:           make(chan *proto.FlowUpdate, channelDepth),
+		sinkChan:           make(chan bucketing.Sink, 10),
 		rolloverFunc:       time.After,
 		bucketsToAggregate: 20,
 		pushIndex:          30,
@@ -195,7 +199,7 @@ func (a *LogAggregator) Run(startTime int64) {
 			a.handleFlowUpdate(upd)
 		case <-rolloverCh:
 			rolloverCh = a.rolloverFunc(a.rollover())
-			a.maybeEmitFlows()
+			a.buckets.EmitFlowCollections(a.sink)
 		case req := <-a.listRequests:
 			req.respCh <- a.queryFlows(req.req)
 		case req := <-a.streamRequests:
@@ -204,13 +208,21 @@ func (a *LogAggregator) Run(startTime int64) {
 			a.backfill(stream, req.req)
 		case id := <-a.streams.closedStreams():
 			a.streams.close(id)
+		case sink := <-a.sinkChan:
+			logrus.WithField("sink", sink).Info("Setting aggregator sink")
+			a.sink = sink
+			a.buckets.EmitFlowCollections(a.sink)
 		case <-a.done:
 			logrus.Warn("Aggregator shutting down")
 			return
 		}
 	}
 }
 
+func (a *LogAggregator) SetSink(s bucketing.Sink) {
+	a.sinkChan <- s
+}
+
 // Receive is used to send a flow update to the aggregator.
 func (a *LogAggregator) Receive(f *proto.FlowUpdate) {
 	timeout := time.After(5 * time.Second)
@@ -222,25 +234,6 @@ func (a *LogAggregator) Receive(f *proto.FlowUpdate) {
 	}
 }
 
-func (a *LogAggregator) maybeEmitFlows() {
-	if a.sink == nil {
-		logrus.Debug("No sink configured, skip flow emission")
-		return
-	}
-
-	flows := a.buckets.FlowCollection()
-	if flows == nil {
-		// We've already pushed this bucket, so we can skip it. We'll emit the next flow once
-		// bucketsToAggregate buckets have been rolled over.
-		logrus.Debug("Delaying flow emission, no new flows to emit")
-		return
-	}
-
-	if len(flows.Flows) > 0 {
-		a.sink.Receive(flows)
-	}
-}
-
 // Stream returns a new Stream from the aggregator. It uses a channel to synchronously request the stream
 // from the aggregator.
 func (a *LogAggregator) Stream(req *proto.FlowStreamRequest) (*Stream, error) {