Refactor ioReaderClient separate logic for --first/--last/--all

Signed-off-by: Chance Zibolski <[email protected]>

Author: chancez

cmd/observe/io_reader_observer.go

@@ -96,46 +96,57 @@ func newIOReaderClient(ctx context.Context, scanner *bufio.Scanner, request *obs
 }
 
 func (c *ioReaderClient) Recv() (*observer.GetFlowsResponse, error) {
-	if c.request.GetNumber() > 0 && c.flowsReturned >= c.request.GetNumber() {
+	if c.returnedEnoughFlows() {
 		return nil, io.EOF
 	}
+
 	for c.scanner.Scan() {
-		line := c.scanner.Text()
-		var res observer.GetFlowsResponse
-		err := protojson.Unmarshal(c.scanner.Bytes(), &res)
-		if err != nil {
-			logger.Logger.WithError(err).WithField("line", line).Warn("Failed to unmarshal json to flow")
-			continue
-		}
-		if c.request.Since != nil && c.request.Since.AsTime().After(res.Time.AsTime()) {
-			continue
-		}
-		if c.request.Until != nil && c.request.Until.AsTime().Before(res.Time.AsTime()) {
-			continue
-		}
-		if !filters.Apply(c.allow, c.deny, &v1.Event{Timestamp: res.Time, Event: res.GetFlow()}) {
+		res, ok := c.unmarshalNext()
+		if !ok {
 			continue
 		}
+
 		switch {
-		case c.buffer != nil && c.request.GetNumber() != math.MaxUint64: // --last
+		case c.isLast():
 			// store flows in a FIFO buffer, effectively keeping the last N flows
 			// until we finish reading from the stream
-			c.buffer.Add(&res)
+			c.buffer.Add(res)
 			continue
-		case c.request.GetFirst() && c.request.GetNumber() != 0 && c.request.GetNumber() != math.MaxUint64: // --first
+		case c.isFirst():
 			// track number of flows returned, so we can exit once we've given back N flows
 			c.flowsReturned++
-			return &res, nil
-		default:
-			return &res, nil
+			return res, nil
+		default: // --all
+			return res, nil
 		}
 	}
+
 	if err := c.scanner.Err(); err != nil {
 		return nil, err
 	}
 
+	if res := c.popFromLastBuffer(); res != nil {
+		return res, nil
+	}
+
+	return nil, io.EOF
+}
+
+func (c *ioReaderClient) isFirst() bool {
+	return c.request.GetFirst() && c.request.GetNumber() != 0 && c.request.GetNumber() != math.MaxUint64
+}
+
+func (c *ioReaderClient) isLast() bool {
+	return c.buffer != nil && c.request.GetNumber() != math.MaxUint64
+}
+
+func (c *ioReaderClient) returnedEnoughFlows() bool {
+	return c.request.GetNumber() > 0 && c.flowsReturned >= c.request.GetNumber()
+}
+
+func (c *ioReaderClient) popFromLastBuffer() *observer.GetFlowsResponse {
 	// Handle --last by iterating over our FIFO and returning one item each time.
-	if c.buffer != nil && c.request.GetNumber() != math.MaxUint64 {
+	if c.isLast() {
 		if len(c.resps) == 0 {
 			// Iterate over the buffer and store them in a slice, because we cannot
 			// index into the ring buffer itself
@@ -150,8 +161,28 @@ func (c *ioReaderClient) Recv() (*observer.GetFlowsResponse, error) {
 		if len(c.resps) > int(c.flowsReturned) {
 			resp := c.resps[c.flowsReturned]
 			c.flowsReturned++
-			return resp, nil
+			return resp
 		}
 	}
-	return nil, io.EOF
+	return nil
+}
+
+func (c *ioReaderClient) unmarshalNext() (*observer.GetFlowsResponse, bool) {
+	line := c.scanner.Text()
+	var res observer.GetFlowsResponse
+	err := protojson.Unmarshal(c.scanner.Bytes(), &res)
+	if err != nil {
+		logger.Logger.WithError(err).WithField("line", line).Warn("Failed to unmarshal json to flow")
+		return nil, false
+	}
+	if c.request.Since != nil && c.request.Since.AsTime().After(res.Time.AsTime()) {
+		return nil, false
+	}
+	if c.request.Until != nil && c.request.Until.AsTime().Before(res.Time.AsTime()) {
+		return nil, false
+	}
+	if !filters.Apply(c.allow, c.deny, &v1.Event{Timestamp: res.Time, Event: res.GetFlow()}) {
+		return nil, false
+	}
+	return &res, true
 }