hubble/recorder: Extend the API to allow stopping a recording automatically

api/v1/recorder/recorder.proto

@@ -15,6 +15,7 @@
 syntax = "proto3";
 
 import "google/protobuf/timestamp.proto";
+import "google/protobuf/duration.proto";
 
 package recorder;
 
@@ -49,6 +50,26 @@ message StartRecording {
     // max_capture_length specifies the maximum packet length.
     // Full packet length will be captured if absent/zero.
     uint32 max_capture_length = 3;
+
+    // stop_condition defines conditions which will cause the recording to
+    // stop early after any of the stop conditions has been hit
+    StopCondition stop_condition = 4;
+}
+
+// StopCondition defines one or more conditions which cause the recording to
+// stop after they have been hit. Stop conditions are ignored if they are
+// absent or zero-valued. If multiple conditions are defined, the recording
+// stops after the first one is hit.
+message StopCondition {
+    // bytes_captured_count stops the recording after at least this many bytes
+    // have been captured. Note: The resulting file might be slightly larger due
+    // to added pcap headers.
+    uint64 bytes_captured_count = 1;
+    // packets_captured_count stops the recording after at least this many packets have
+    // been captured.
+    uint64 packets_captured_count = 2;
+    // time_elapsed stops the recording after this duration has elapsed.
+    google.protobuf.Duration time_elapsed = 3;
 }
 
 // FileSinkConfiguration configures the file output. Possible future additions

pkg/hubble/recorder/service.go

@@ -342,20 +342,26 @@ func (s *Service) startRecording(
 	})
 	scopedLog.Debug("starting new recording")
 
+	stop := req.GetStopCondition()
 	config := sink.PcapSink{
 		RuleID: ruleID,
 		Header: pcap.Header{
 			SnapshotLength: capLen,
 			Datalink:       pcap.Ethernet,
 		},
 		Writer: pcap.NewWriter(f),
+		StopCondition: sink.StopConditions{
+			PacketsCaptured: stop.GetPacketsCapturedCount(),
+			BytesCaptured:   stop.GetBytesCapturedCount(),
+			DurationElapsed: stop.GetTimeElapsed().AsDuration(),
+		},
 	}
 
-	handle, err = s.dispatch.StartSink(ctx, config)
-	if err != nil {
-		return nil, "", err
-	}
-
+	// Upserting a new recorder can take up to a few seconds due to datapath
+	// regeneration. To avoid having the stop condition timer on the sink
+	// already running while the recorder is still being upserted, we install
+	// the recorder before the sink. This is safe, as sink.Dispatch silently
+	// ignores recordings for unknown sinks.
 	recInfo := &recorder.RecInfo{
 		ID:      recorder.ID(ruleID),
 		CapLen:  uint16(capLen),
@@ -366,6 +372,11 @@ func (s *Service) startRecording(
 		return nil, "", err
 	}
 
+	handle, err = s.dispatch.StartSink(ctx, config)
+	if err != nil {
+		return nil, "", err
+	}
+
 	// Ensure to delete the above recorder when the sink has stopped
 	go func() {
 		<-handle.Done

pkg/hubble/recorder/sink/dispatch.go

@@ -78,11 +78,20 @@ type Statistics struct {
 	BytesLost      uint64
 }
 
+// StopConditions defines a set of values which cause the sink to stop
+// recording if any of them are hit. Zero-valued conditions are ignored.
+type StopConditions struct {
+	PacketsCaptured uint64
+	BytesCaptured   uint64
+	DurationElapsed time.Duration
+}
+
 // PcapSink defines the parameters of a sink which writes to a pcap.RecordWriter
 type PcapSink struct {
-	RuleID uint16
-	Header pcap.Header
-	Writer pcap.RecordWriter
+	RuleID        uint16
+	Header        pcap.Header
+	Writer        pcap.RecordWriter
+	StopCondition StopConditions
 }
 
 // Dispatch implements consumer.MonitorConsumer and dispatches incoming
@@ -121,6 +130,7 @@ func NewDispatch(sinkQueueSize int) (*Dispatch, error) {
 // The sink is unregistered automatically when it stops. A sink is stopped for
 // one of the following four reasons. In all cases, Handle.Done will be closed.
 //  - Explicitly via Handle.Stop (Handle.Err() == nil)
+//  - When one of the p.StopCondition is hit (Handle.Err() == nil)
 //  - When the context ctx is cancelled (Handle.Err() != nil)
 //  - When an error occurred (Handle.Err() != nil)
 func (d *Dispatch) StartSink(ctx context.Context, p PcapSink) (*Handle, error) {

pkg/hubble/recorder/sink/sink.go

@@ -16,6 +16,7 @@ package sink
 
 import (
 	"context"
+	"time"
 
 	"github.com/cilium/cilium/pkg/hubble/recorder/pcap"
 	"github.com/cilium/cilium/pkg/lock"
@@ -39,6 +40,7 @@ type sink struct {
 // startSink creates a queue and go routine for the sink. The spawned go
 // routine will run until one of the following happens:
 //  - sink.stop is called
+//  - a p.StopCondition is reached
 //  - ctx is cancelled
 //  - an error occurred
 func startSink(ctx context.Context, p PcapSink, queueSize int) *sink {
@@ -69,6 +71,16 @@ func startSink(ctx context.Context, p PcapSink, queueSize int) *sink {
 			close(s.done)
 		}()
 
+		stop := p.StopCondition
+		var stopAfter <-chan time.Time
+		if stop.DurationElapsed != 0 {
+			stopTimer := time.NewTimer(stop.DurationElapsed)
+			defer func() {
+				stopTimer.Stop()
+			}()
+			stopAfter = stopTimer.C
+		}
+
 		if err = p.Writer.WriteHeader(p.Header); err != nil {
 			return
 		}
@@ -87,12 +99,19 @@ func startSink(ctx context.Context, p PcapSink, queueSize int) *sink {
 					return
 				}
 
-				s.addToStatistics(Statistics{
+				stats := s.addToStatistics(Statistics{
 					PacketsWritten: 1,
 					BytesWritten:   uint64(rec.inclLen),
 				})
+				if (stop.PacketsCaptured > 0 && stats.PacketsWritten >= stop.PacketsCaptured) ||
+					(stop.BytesCaptured > 0 && stats.BytesWritten >= stop.BytesCaptured) {
+					return
+				}
 			case <-s.shutdown:
 				return
+			case <-stopAfter:
+				// duration of stop condition has been reached
+				return
 			case <-ctx.Done():
 				err = ctx.Err()
 				return
@@ -108,19 +127,24 @@ func (s *sink) stop() {
 	close(s.shutdown)
 }
 
-func (s *sink) addToStatistics(add Statistics) {
+// addToStatistics adds add to the current statistics and returns the resulting
+// value.
+func (s *sink) addToStatistics(add Statistics) (result Statistics) {
 	s.mutex.Lock()
 	s.stats.BytesWritten += add.BytesWritten
 	s.stats.PacketsWritten += add.PacketsWritten
 	s.stats.BytesLost += add.BytesLost
 	s.stats.PacketsLost += add.PacketsLost
+	result = s.stats
 	s.mutex.Unlock()
 
 	// non-blocking send
 	select {
 	case s.trigger <- struct{}{}:
 	default:
 	}
+
+	return result
 }
 
 // enqueue submits a new record to this sink. If the sink is not keeping up,