larksuite-cli/internal/event/consume/loop.go

// Copyright (c) 2026 Lark Technologies Pte. Ltd.
// SPDX-License-Identifier: MIT

package consume

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io/fs"
	"net"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/itchyny/gojq"
	"github.com/larksuite/cli/internal/event"
	"github.com/larksuite/cli/internal/event/protocol"
)

// consumeLoop reads events and dispatches to workers; cancels on terminal sink errors.
func consumeLoop(ctx context.Context, conn net.Conn, br *bufio.Reader, keyDef *event.KeyDefinition, opts Options, subscriptionID string, lastForKey *bool, emitted *atomic.Int64) error {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	sink, err := newSink(opts)
	if err != nil {
		return err
	}

	// Compile before worker goroutines start to avoid a data race on jqCode.
	var jqCode *gojq.Code
	if opts.JQExpr != "" {
		jqCode, err = CompileJQ(opts.JQExpr)
		if err != nil {
			return err
		}
	}

	bufSize := keyDef.BufferSize
	if bufSize <= 0 {
		bufSize = event.DefaultBufferSize
	}
	socketCh := make(chan *protocol.Event, bufSize)

	// stopReader lets shutdown preempt the reader so PreShutdownCheck can reuse conn.
	stopReader := make(chan struct{})
	readerDone := make(chan struct{})

	// ReadBytes (not Scanner) so mid-frame read deadlines don't drop buffered bytes.
	go func() {
		defer close(readerDone)
		defer close(socketCh)
		var buf []byte
		var lastSeq uint64 // per-conn monotonic; gaps = bus drop-oldest backpressure
		for {
			select {
			case <-stopReader:
				return
			default:
			}
			conn.SetReadDeadline(time.Now().Add(200 * time.Millisecond))
			chunk, err := br.ReadBytes('\n')
			if len(chunk) > 0 {
				// Cap accumulator: dribbling multi-MB lines past 200ms deadlines could grow buf unbounded.
				if len(buf)+len(chunk) > protocol.MaxFrameBytes {
					if !opts.Quiet {
						fmt.Fprintf(opts.ErrOut,
							"WARN: dropping oversized frame (>%d bytes) from bus\n", protocol.MaxFrameBytes)
					}
					buf = nil
					continue
				}
				buf = append(buf, chunk...)
			}
			if err != nil {
				var ne net.Error
				if errors.As(err, &ne) && ne.Timeout() {
					continue
				}
				return
			}
			line := buf
			if n := len(line); n > 0 && line[n-1] == '\n' {
				line = line[:n-1]
			}
			buf = nil

			msg, decErr := protocol.Decode(line)
			if decErr != nil {
				continue
			}
			switch m := msg.(type) {
			case *protocol.Event:
				if lastSeq > 0 && m.Seq > 0 && m.Seq > lastSeq+1 {
					gap := m.Seq - lastSeq - 1
					if !opts.Quiet {
						fmt.Fprintf(opts.ErrOut,
							"WARN: event seq gap %d->%d, missed %d events (dropped by bus backpressure)\n",
							lastSeq, m.Seq, gap)
					}
				}
				// Only advance forward — concurrent publishers can deliver out-of-order.
				if m.Seq > lastSeq {
					lastSeq = m.Seq
				}
				select {
				case socketCh <- m:
				default:
					// drop-oldest back-pressure
					select {
					case <-socketCh:
					default:
					}
					select {
					case socketCh <- m:
					default:
					}
					if !opts.Quiet {
						fmt.Fprintf(opts.ErrOut, "WARN: consume backpressure, dropped oldest event\n")
					}
				}
			case *protocol.SourceStatus:
				if !opts.Quiet {
					if m.Detail != "" {
						fmt.Fprintf(opts.ErrOut, "[source] %s: %s (%s)\n", m.Source, m.State, m.Detail)
					} else {
						fmt.Fprintf(opts.ErrOut, "[source] %s: %s\n", m.Source, m.State)
					}
				}
			default:
				// forward-compatible: ignore unknown message types
			}
		}
	}()

	workers := keyDef.Workers
	if workers <= 0 {
		workers = 1
	}

	var wg sync.WaitGroup
	wg.Add(workers)
	for i := 0; i < workers; i++ {
		go func() {
			defer wg.Done()
			for evt := range socketCh {
				wrote, err := processAndOutput(ctx, keyDef, evt, opts, sink, jqCode)
				if wrote {
					emitted.Add(1)
					// cancel inner ctx so shutdown goes through normal cleanup, not conn rip.
					if checkMaxEvents(opts, emitted) {
						cancel()
						return
					}
				}
				if err != nil {
					if isTerminalSinkError(err) {
						if !opts.Quiet {
							fmt.Fprintf(opts.ErrOut, "consume: output pipe closed (%v), shutting down\n", err)
						}
						cancel()
						return
					}
					if !opts.Quiet {
						fmt.Fprintf(opts.ErrOut, "WARN: sink write failed, skipping event: %v\n", err)
					}
				}
			}
		}()
	}

	allDone := make(chan struct{})
	go func() {
		wg.Wait()
		close(allDone)
	}()

	select {
	case <-ctx.Done():
		// Drain reader so PreShutdownCheck has exclusive conn.
		close(stopReader)
		<-readerDone
		conn.SetReadDeadline(time.Time{})
		*lastForKey = checkLastForKey(conn, opts.EventKey, subscriptionID)
		conn.Close()
	case <-allDone:
		// bus-side close; can't query, assume last
		*lastForKey = true
	}

	wg.Wait()

	return nil
}

// processAndOutput returns (wrote, err); err non-nil only for sink.Write failures.
func processAndOutput(ctx context.Context, keyDef *event.KeyDefinition, evt *protocol.Event, opts Options, sink Sink, jqCode *gojq.Code) (bool, error) {
	raw := &event.RawEvent{
		EventType: evt.EventType,
		Payload:   evt.Payload,
	}

	// Synchronous Match filter runs before any work (Process / sink write).
	if keyDef.Match != nil && !keyDef.Match(raw, opts.Params) {
		return false, nil
	}

	var result json.RawMessage

	if keyDef.Process != nil {
		var err error
		result, err = keyDef.Process(ctx, opts.Runtime, raw, opts.Params)
		if err != nil {
			if !opts.Quiet {
				fmt.Fprintf(opts.ErrOut, "WARN: Process error: %v\n", err)
			}
			return false, nil
		}
		if result == nil {
			return false, nil
		}
	} else {
		result = evt.Payload
	}

	if jqCode != nil {
		filtered, err := applyJQ(jqCode, result)
		if err != nil {
			if !opts.Quiet {
				fmt.Fprintf(opts.ErrOut, "WARN: JQ error: %v\n", err)
			}
			return false, nil
		}
		if filtered == nil {
			return false, nil
		}
		result = filtered
	}

	if err := sink.Write(result); err != nil {
		return false, err
	}
	return true, nil
}

// isTerminalSinkError reports if the output channel is permanently broken (EPIPE/ErrClosed).
func isTerminalSinkError(err error) bool {
	if err == nil {
		return false
	}
	if errors.Is(err, syscall.EPIPE) {
		return true
	}
	if errors.Is(err, fs.ErrClosed) {
		return true
	}
	return false
}