larksuite--cli/internal/event/bus/conn.go

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

package bus

import (
	"bufio"
	"bytes"
	"log"
	"net"
	"sync"
	"sync/atomic"
	"time"

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

const (
	sendChCap    = 100
	writeTimeout = 5 * time.Second
)

// Conn represents a single consume client connection in the Bus.
type Conn struct {
	conn            net.Conn
	reader          *bufio.Reader
	sendCh          chan interface{}
	sendMu          sync.Mutex // serialises drop+push atomically
	writeMu         sync.Mutex // serialises all net.Conn writes (Encode+SetWriteDeadline is a 2-call sequence)
	eventKey        string
	eventTypes      []string
	subID           string
	pid             int
	onClose         func(*Conn)
	checkLastForKey func(scope string) bool
	logger          *log.Logger
	closed          chan struct{}
	closeOnce       sync.Once
	received        atomic.Int64  // events fanned out to us (post-filter)
	seqCounter      atomic.Uint64 // per-conn monotonic seq assigned by Hub.Publish
	dropped         atomic.Int64  // events evicted via drop-oldest backpressure
}

// NewConn creates a Conn; pass a reader with pre-buffered bytes (handoff from Bus.handleConn) or nil for a fresh one.
func NewConn(conn net.Conn, reader *bufio.Reader, eventKey string, eventTypes []string, pid int, subID string) *Conn {
	if reader == nil {
		reader = bufio.NewReader(conn)
	}
	return &Conn{
		conn:       conn,
		reader:     reader,
		sendCh:     make(chan interface{}, sendChCap),
		eventKey:   eventKey,
		eventTypes: eventTypes,
		pid:        pid,
		subID:      subID,
		closed:     make(chan struct{}),
	}
}

// SubscriptionID returns the subscription identity. Falls back to EventKey
// when the stored subID is empty (legacy clients / no-SubscriptionKey EventKeys).
func (c *Conn) SubscriptionID() string {
	if c.subID == "" {
		return c.eventKey
	}
	return c.subID
}

func (c *Conn) SetOnClose(fn func(*Conn)) { c.onClose = fn }

// SetCheckLastForKey: returning true means "you are the last subscriber, run cleanup".
func (c *Conn) SetCheckLastForKey(fn func(string) bool) { c.checkLastForKey = fn }

// SetLogger attaches a logger (nil tolerated).
func (c *Conn) SetLogger(l *log.Logger) { c.logger = l }

func (c *Conn) EventKey() string         { return c.eventKey }
func (c *Conn) EventTypes() []string     { return c.eventTypes }
func (c *Conn) SendCh() chan interface{} { return c.sendCh }
func (c *Conn) PID() int                 { return c.pid }
func (c *Conn) IncrementReceived()       { c.received.Add(1) }
func (c *Conn) Received() int64          { return c.received.Load() }

// NextSeq returns the next monotonic seq for this conn (first call returns 1).
func (c *Conn) NextSeq() uint64 { return c.seqCounter.Add(1) }

func (c *Conn) DroppedCount() int64 { return c.dropped.Load() }
func (c *Conn) IncrementDropped()   { c.dropped.Add(1) }

// Start launches the sender and reader goroutines; call exactly once.
func (c *Conn) Start() {
	go c.SenderLoop()
	go c.ReaderLoop()
}

// writeFrame is the sole write path, serialised via writeMu.
func (c *Conn) writeFrame(msg interface{}) error {
	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	if err := c.conn.SetWriteDeadline(time.Now().Add(writeTimeout)); err != nil {
		return err
	}
	return protocol.Encode(c.conn, msg)
}

// SenderLoop exits on closed (not sendCh close) so Hub.Publish can send without panic risk.
func (c *Conn) SenderLoop() {
	for {
		select {
		case <-c.closed:
			return
		case msg := <-c.sendCh:
			if err := c.writeFrame(msg); err != nil {
				if c.logger != nil {
					c.logger.Printf("WARN: write to pid=%d failed: %v", c.pid, err)
				}
				c.shutdown()
				return
			}
		}
	}
}

// ReaderLoop reads control messages (Bye, PreShutdownCheck) until EOF.
func (c *Conn) ReaderLoop() {
	for {
		line, err := protocol.ReadFrame(c.reader)
		if err != nil {
			break
		}
		line = bytes.TrimRight(line, "\n")
		if len(line) == 0 {
			continue
		}
		msg, err := protocol.Decode(line)
		if err != nil {
			continue
		}
		c.handleControlMessage(msg)
	}
	c.shutdown()
}

func (c *Conn) handleControlMessage(msg interface{}) {
	switch msg.(type) {
	case *protocol.Bye:
		c.shutdown()
	case *protocol.PreShutdownCheck:
		// Use the connection's own authoritative subscription identity rather
		// than recomputing from the incoming message: a stale or mismatched
		// PreShutdownCheck must not ask about the wrong scope (which would
		// suppress or mistrigger per-subscription cleanup). Conn.SubscriptionID()
		// already falls back to EventKey when its stored subID is empty.
		scope := c.SubscriptionID()
		lastForKey := true
		if c.checkLastForKey != nil {
			lastForKey = c.checkLastForKey(scope)
		}
		ack := protocol.NewPreShutdownAck(lastForKey)
		if err := c.writeFrame(ack); err != nil && c.logger != nil {
			c.logger.Printf("WARN: pre_shutdown_ack to pid=%d failed: %v", c.pid, err)
		}
	}
}

func (c *Conn) shutdown() {
	c.closeOnce.Do(func() {
		close(c.closed)
		c.conn.Close()
		// sendCh is NOT closed: would race with Hub.Publish holding SendCh() after RUnlock.
		if c.onClose != nil {
			c.onClose(c)
		}
	})
}

// TrySend enqueues non-evictively under sendMu so it respects PushDropOldest's atomicity contract.
func (c *Conn) TrySend(msg interface{}) bool {
	c.sendMu.Lock()
	defer c.sendMu.Unlock()
	select {
	case c.sendCh <- msg:
		return true
	default:
		return false
	}
}

// PushDropOldest enqueues msg; on full channel evicts one oldest and retries, atomically under sendMu.
// Returns (enqueued, dropped). A rare concurrent drain may make drop unnecessary — still succeeds with dropped=false.
func (c *Conn) PushDropOldest(msg interface{}) (enqueued, dropped bool) {
	c.sendMu.Lock()
	defer c.sendMu.Unlock()
	select {
	case c.sendCh <- msg:
		return true, false
	default:
	}
	select {
	case <-c.sendCh:
		dropped = true
	default:
	}
	select {
	case c.sendCh <- msg:
		return true, dropped
	default:
		return false, dropped
	}
}

// Close is idempotent.
func (c *Conn) Close() {
	c.shutdown()
}