AI-tianyan/internal/infer/client.go

package infer

import (
	"context"
	"crypto/rand"
	"encoding/base64"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"sync"
	"time"

	"tianyan-edge/internal/config"
	"tianyan-edge/internal/event"
	"tianyan-edge/internal/stream"
)

type frameMsg struct {
	StreamID int     `json:"stream_id"`
	DeviceID string  `json:"device_id"`
	URL      string  `json:"url"`
	TS       float64 `json:"ts"`
	JPEGB64  string  `json:"jpeg_b64"`
}

type detection struct {
	Class string    `json:"class"`
	Conf  float64   `json:"conf"`
	BBox  []float64 `json:"bbox"`
}

type resultMsg struct {
	StreamID   int         `json:"stream_id"`
	DeviceID   string      `json:"device_id"`
	TS         float64     `json:"ts"`
	Detections []detection `json:"detections"`
}

type dedupState struct {
	mu   sync.Mutex
	seen map[string]float64
}

type Client struct {
	cfg       *config.Config
	frames    <-chan stream.Frame
	events    chan<- event.SuspectedEvent
	dedup     *dedupState
	updates   <-chan *config.Config
}

func NewClient(cfg *config.Config, frames <-chan stream.Frame, events chan<- event.SuspectedEvent, updates <-chan *config.Config) *Client {
	return &Client{
		cfg:     cfg,
		frames:  frames,
		events:  events,
		dedup:   &dedupState{seen: make(map[string]float64)},
		updates: updates,
	}
}

func (c *Client) Run(ctx context.Context) {
	var wg sync.WaitGroup

	// Worker management loop
	numWorkers := c.cfg.InferWorkers
	activeWorkers := make([]context.CancelFunc, numWorkers)

	for i := 0; i < numWorkers; i++ {
		workerCtx, cancel := context.WithCancel(ctx)
		activeWorkers[i] = cancel
		wg.Add(1)
		go func(id int, wCtx context.Context) {
			defer wg.Done()
			c.workerLoop(wCtx, id)
		}(i, workerCtx)
	}

	for {
		select {
		case <-ctx.Done():
			for _, cancel := range activeWorkers {
				cancel()
			}
			wg.Wait()
			return
		case newCfg := <-c.updates:
			c.cfg = newCfg
			target := c.cfg.InferWorkers
			if target < 0 {
				target = 0
			}

			// Scale up
			for i := len(activeWorkers); i < target; i++ {
				workerCtx, cancel := context.WithCancel(ctx)
				activeWorkers = append(activeWorkers, cancel)
				wg.Add(1)
				go func(id int, wCtx context.Context) {
					defer wg.Done()
					c.workerLoop(wCtx, id)
				}(i, workerCtx)
			}
			// Scale down
			if target < len(activeWorkers) {
				for i := target; i < len(activeWorkers); i++ {
					activeWorkers[i]()
				}
				activeWorkers = activeWorkers[:target]
			}
			log.Printf("infer: workers scaled to %d", target)
		}
	}
}

func (c *Client) workerLoop(ctx context.Context, workerID int) {
	for {
		select {
		case <-ctx.Done():
			return
		default:
		}
		conn, err := net.Dial("unix", c.cfg.InferSocket)
		if err != nil {
			log.Printf("infer[w%d]: socket connect failed: %v, retry 2s", workerID, err)
			time.Sleep(2 * time.Second)
			continue
		}
		log.Printf("infer[w%d]: socket connected", workerID)
		if !c.loop(ctx, conn, workerID) {
			conn.Close()
			time.Sleep(1 * time.Second)
		}
		conn.Close()
	}
}

func (c *Client) loop(ctx context.Context, conn net.Conn, workerID int) bool {
	for {
		select {
		case <-ctx.Done():
			return false
		case f := <-c.frames:
			if err := c.send(conn, f); err != nil {
				return true
			}
			res, err := c.recv(conn)
			if err != nil {
				return true
			}
			c.emitEvents(f, res)
		}
	}
}

func (c *Client) emitEvents(f stream.Frame, res *resultMsg) {
	for _, d := range res.Detections {
		if c.isDuplicate(f.DeviceID, d.Class, f.TS) {
			continue
		}
		ev := event.SuspectedEvent{
			EventID:   mustUUID(),
			EdgeID:    c.cfg.EdgeID,
			DeviceID:  f.DeviceID,
			StreamURL: f.URL,
			TS:        f.TS,
			Class:     d.Class,
			Conf:      d.Conf,
			BBox:      d.BBox,
			ImageB64:  base64.StdEncoding.EncodeToString(f.JPEG),
		}
		select {
		case c.events <- ev:
		default:
			log.Println("infer: event queue full, dropping")
		}
	}
}

func (c *Client) isDuplicate(deviceID, className string, ts float64) bool {
	window := float64(c.cfg.DedupWindowSec)
	key := fmt.Sprintf("%s|%s", deviceID, className)
	c.dedup.mu.Lock()
	defer c.dedup.mu.Unlock()
	last, ok := c.dedup.seen[key]
	if ok && ts-last < window {
		return true
	}
	c.dedup.seen[key] = ts
	if len(c.dedup.seen) > 5000 {
		for k, v := range c.dedup.seen {
			if ts-v > 2*window {
				delete(c.dedup.seen, k)
			}
		}
	}
	return false
}

func (c *Client) send(conn net.Conn, f stream.Frame) error {
	msg := frameMsg{StreamID: f.StreamID, DeviceID: f.DeviceID, URL: f.URL,
		TS: f.TS, JPEGB64: base64.StdEncoding.EncodeToString(f.JPEG)}
	data, err := json.Marshal(msg)
	if err != nil {
		return err
	}
	conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
	defer conn.SetWriteDeadline(time.Time{})
	return writeMsg(conn, data)
}

func (c *Client) recv(conn net.Conn) (*resultMsg, error) {
	// Set read deadline to prevent blocking forever if inference server hangs
	// 30s timeout covers even slow NPU inference on large images
	conn.SetReadDeadline(time.Now().Add(30 * time.Second))
	defer conn.SetReadDeadline(time.Time{}) // reset after read

	data, err := readMsg(conn)
	if err != nil {
		return nil, err
	}
	var r resultMsg
	return &r, json.Unmarshal(data, &r)
}

func writeMsg(conn net.Conn, data []byte) error {
	hdr := make([]byte, 4)
	binary.BigEndian.PutUint32(hdr, uint32(len(data)))
	if _, err := conn.Write(hdr); err != nil {
		return err
	}
	_, err := conn.Write(data)
	return err
}

func readMsg(conn net.Conn) ([]byte, error) {
	hdr := make([]byte, 4)
	if _, err := io.ReadFull(conn, hdr); err != nil {
		return nil, err
	}
	size := binary.BigEndian.Uint32(hdr)
	buf := make([]byte, size)
	_, err := io.ReadFull(conn, buf)
	return buf, err
}

func mustUUID() string {
	b := make([]byte, 16)
	_, _ = rand.Read(b)
	return hex.EncodeToString(b[:4]) + "-" + hex.EncodeToString(b[4:6]) + "-" +
		hex.EncodeToString(b[6:8]) + "-" + hex.EncodeToString(b[8:10]) + "-" +
		hex.EncodeToString(b[10:])
}