oko_public/terp/main.go

package main

import (
	"context"
	"flag"
	"log"
	"net"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"frame"
)

const (
	workerCount = 4

	// connLimit bounds concurrent inbound connections so a connection flood
	// cannot exhaust goroutines or file descriptors.
	connLimit = 16

	// clipChanCap gives in-flight ingest handlers somewhere to park a decoded
	// clip without holding the connection's socket open for the workers.
	clipChanCap = 32

	shutdownGrace = 5 * time.Second
)

func main() {
	var (
		listenAddr     string
		sendoffAddr    string
		classifierGlob string
		motionPercent  int
		passphrase     string
		authToken      string
		pbkdf2Iter     int
	)

	flag.StringVar(&listenAddr, "listen", ":8081", "address to listen for incoming clips")
	flag.StringVar(&sendoffAddr, "sendoff", "localhost:8082", "address to forward classified clips")
	flag.StringVar(&classifierGlob, "classifiers", "classifiers/*.xml", "glob pattern for Haar cascade XMLs")
	flag.IntVar(&motionPercent, "motion-percent", 0, "minimum highest motion percent [0-100] required to classify and forward (0 = always classify)")
	flag.StringVar(&passphrase, "passphrase", "", "decryption passphrase (empty = no encryption)")
	flag.StringVar(&authToken, "auth-token", "", "shared pipeline auth token (or OKO_AUTH_TOKEN)")
	flag.IntVar(&pbkdf2Iter, "pbkdf2-iters", frame.DefaultPBKDF2Iter, "PBKDF2 iterations for clip encryption (1..1000000)")
	flag.Parse()

	if pbkdf2Iter < 1 || pbkdf2Iter > frame.MaxPBKDF2Iter {
		log.Fatalf("pbkdf2-iters must be in [1, %d], got %d", frame.MaxPBKDF2Iter, pbkdf2Iter)
	}
	if motionPercent < 0 || motionPercent > 100 {
		log.Fatalf("motion-percent must be in [0, 100], got %d", motionPercent)
	}
	frame.SetDefaultPBKDF2Iter(pbkdf2Iter)

	tok := authToken
	if tok == "" {
		tok = os.Getenv("OKO_AUTH_TOKEN")
	}
	if tok == "" {
		log.Fatal("authentication required: set -auth-token or OKO_AUTH_TOKEN")
	}
	frame.SetAuthToken(tok)

	// Load every classifier once at startup instead of per frame.
	classifiers, err := frame.LoadClassifiers([]string{classifierGlob})
	if err != nil {
		log.Fatalf("LoadClassifiers: %v", err)
	}
	defer func() {
		for i := range classifiers {
			classifiers[i].Close()
		}
	}()
	log.Printf("Loaded %d classifier(s)", len(classifiers))

	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer cancel()

	clipChan := make(chan frame.Clip, clipChanCap)
	sendoffChan := make(chan frame.Clip, 8)
	connSem := make(chan struct{}, connLimit)

	ln, err := net.Listen("tcp", listenAddr)
	if err != nil {
		log.Fatalf("Listen %s: %v", listenAddr, err)
	}
	log.Printf("Listening on %s", listenAddr)

	// Accept loop
	var acceptWg sync.WaitGroup
	rejected := frame.NewMalformedAttempts()
	go func() {
		<-ctx.Done()
		ln.Close()
	}()
	go func() {
		for {
			conn, err := ln.Accept()
			if err != nil {
				if ctx.Err() != nil {
					return
				}
				log.Printf("Accept error: %v", err)
				continue
			}
			select {
			case connSem <- struct{}{}:
			default:
				log.Printf("Rejecting connection from %s: too many concurrent connections", conn.RemoteAddr())
				conn.Close()
				continue
			}
			acceptWg.Add(1)
			go func(c net.Conn) {
				defer func() { <-connSem }()
				defer acceptWg.Done()
				frame.HandleClipConn(c, rejected, func(clp frame.Clip) {
					clipChan <- clp
				})
			}(conn)
		}
	}()

	// Bounded workers: motion-gate → classify → sendoffChan
	var classifyWg sync.WaitGroup
	classifyWg.Add(workerCount)
	for w := 0; w < workerCount; w++ {
		go func() {
			defer classifyWg.Done()
			for clp := range clipChan {
				classifyClip(clp, sendoffChan, classifiers, motionPercent, passphrase)
			}
		}()
	}

	// Sendoff: forward processed clips to the coordinator with retry
	var sendoffWg sync.WaitGroup
	sendoffWg.Add(1)
	go func() {
		defer sendoffWg.Done()
		for clp := range sendoffChan {
			if err := clp.SendClipRetry(sendoffAddr, 6); err != nil {
				log.Printf("Sendoff to %s failed, clip dropped: %v", sendoffAddr, err)
			}
		}
	}()

	<-ctx.Done()
	log.Println("Shutting down...")

	ln.Close()
	if !frame.WaitGroupTimeout(&acceptWg, shutdownGrace) {
		log.Println("Inbound handlers did not drain; closing channel anyway")
	}
	close(clipChan)
	if !frame.WaitGroupTimeout(&classifyWg, shutdownGrace) {
		log.Println("Classifiers still busy; closing sendoff anyway")
	}
	close(sendoffChan)
	if !frame.WaitGroupTimeout(&sendoffWg, shutdownGrace) {
		log.Println("Sendoff still busy (coordinator retry?); exiting anyway")
	}
	if n := rejected.Total(); n > 0 {
		log.Printf("Rejected %d malformed clips from the network", n)
	}
	log.Println("Exiting")
}

// classifyClip decrypts, motion-gates (skipping recompute when mofin already
// did it), classifies, re-encrypts, and forwards the clip.
func classifyClip(clp frame.Clip, sendoff chan<- frame.Clip, classifiers []frame.Classifier, motionPercent int, passphrase string) {
	if passphrase != "" {
		if err := clp.Decrypt(passphrase); err != nil {
			log.Printf("Decrypt error: %v", err)
			return
		}
	}

	if len(clp.Comparisons) == 0 {
		if err := clp.CountChangedPixels(30); err != nil {
			log.Printf("CountChangedPixels: %v", err)
			return
		}
	}
	highest := clp.GetHighestMotion()
	if motionPercent > 0 && highest < motionPercent {
		log.Printf("Dropped clip with motion=%d%% (below %d%% threshold)", highest, motionPercent)
		return
	}

	frames := clp.Sublimate()
	clp.Detections = make([]map[string][]frame.Detection, len(frames))
	for i, frm := range frames {
		frm.DetectClassifiers(classifiers)
		clp.Detections[i] = frm.Detections
	}

	if passphrase != "" {
		if err := clp.Encrypt(passphrase); err != nil {
			log.Printf("Encrypt error: %v", err)
			return
		}
	}

	hasDetection := false
	for _, dets := range clp.Detections {
		if len(dets) > 0 {
			hasDetection = true
			break
		}
	}

	log.Printf("Forwarding clip with motion=%d%% (detections: %v)", highest, hasDetection)
	sendoff <- clp
}