468 lines
12 KiB
Go
468 lines
12 KiB
Go
package batchstore
|
||
|
||
import (
|
||
"database/sql"
|
||
"encoding/json"
|
||
"fmt"
|
||
"os"
|
||
"path/filepath"
|
||
"sync"
|
||
"time"
|
||
|
||
_ "modernc.org/sqlite"
|
||
)
|
||
|
||
type LogFunc func(string)
|
||
|
||
type Store struct {
|
||
dbPath string
|
||
stateFile string
|
||
cameraName string
|
||
objectLabel string
|
||
cutoffTime string
|
||
|
||
batchTimeout time.Duration
|
||
ignoreBatchLabelTimeout time.Duration
|
||
minObjectPerBatch int
|
||
minDurationPerBatch int
|
||
carryIDs int
|
||
|
||
log LogFunc
|
||
|
||
db *sql.DB
|
||
mu sync.Mutex
|
||
currentState *batchState
|
||
previousState *batchState
|
||
|
||
batchTimer *time.Timer
|
||
ignoreBatchLabel bool
|
||
ignoreBatchLabelTimer *time.Timer
|
||
|
||
shutdownCh chan struct{}
|
||
stopped bool
|
||
}
|
||
|
||
type batchState struct {
|
||
CountingDate string `json:"counting_date"`
|
||
BatchNumber int `json:"batch_number"`
|
||
Count int `json:"count"`
|
||
StartTime string `json:"start_time"`
|
||
LastDetection string `json:"last_detection_time"`
|
||
CountedEventIDs []string `json:"counted_event_ids"`
|
||
}
|
||
|
||
type Config struct {
|
||
DBPath string
|
||
StateFile string
|
||
CameraName string
|
||
ObjectLabel string
|
||
CutoffTime string
|
||
BatchTimeoutSec float64
|
||
IgnoreBatchLabelTimeout float64
|
||
MinObjectPerBatch int
|
||
MinDurationPerBatch int
|
||
}
|
||
|
||
func New(cfg Config, logFn LogFunc) (*Store, error) {
|
||
if logFn == nil {
|
||
logFn = func(s string) { fmt.Println(s) }
|
||
}
|
||
|
||
s := &Store{
|
||
dbPath: cfg.DBPath,
|
||
stateFile: cfg.StateFile,
|
||
cameraName: cfg.CameraName,
|
||
objectLabel: cfg.ObjectLabel,
|
||
cutoffTime: cfg.CutoffTime,
|
||
batchTimeout: time.Duration(cfg.BatchTimeoutSec * float64(time.Second)),
|
||
ignoreBatchLabelTimeout: time.Duration(cfg.IgnoreBatchLabelTimeout * float64(time.Second)),
|
||
minObjectPerBatch: cfg.MinObjectPerBatch,
|
||
minDurationPerBatch: cfg.MinDurationPerBatch,
|
||
carryIDs: 50,
|
||
log: logFn,
|
||
shutdownCh: make(chan struct{}),
|
||
}
|
||
|
||
if err := os.MkdirAll(filepath.Dir(cfg.DBPath), 0755); err != nil {
|
||
return nil, err
|
||
}
|
||
if err := os.MkdirAll(filepath.Dir(cfg.StateFile), 0755); err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
var err error
|
||
s.db, err = sql.Open("sqlite", cfg.DBPath+"?_journal_mode=WAL&_synchronous=NORMAL")
|
||
if err != nil {
|
||
return nil, fmt.Errorf("open db: %w", err)
|
||
}
|
||
|
||
if err := s.initDB(); err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
s.currentState = s.loadState()
|
||
s.previousState = s.currentState
|
||
|
||
return s, nil
|
||
}
|
||
|
||
func (s *Store) initDB() error {
|
||
_, err := s.db.Exec(`
|
||
CREATE TABLE IF NOT EXISTS batches (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
counting_date TEXT NOT NULL,
|
||
batch_number INTEGER NOT NULL,
|
||
camera_name TEXT NOT NULL,
|
||
object_label TEXT NOT NULL,
|
||
count INTEGER NOT NULL,
|
||
start_time TEXT NOT NULL,
|
||
end_time TEXT NOT NULL,
|
||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||
UNIQUE(counting_date, batch_number, camera_name, object_label)
|
||
)
|
||
`)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
_, err = s.db.Exec(`
|
||
CREATE TABLE IF NOT EXISTS daily_summaries (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
counting_date TEXT NOT NULL,
|
||
camera_name TEXT NOT NULL,
|
||
object_label TEXT NOT NULL,
|
||
total_count INTEGER NOT NULL DEFAULT 0,
|
||
total_batches INTEGER NOT NULL DEFAULT 0,
|
||
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||
UNIQUE(counting_date, camera_name, object_label)
|
||
)
|
||
`)
|
||
return err
|
||
}
|
||
|
||
func (s *Store) countingDate(t time.Time) string {
|
||
cutoff, _ := time.Parse("15:04", s.cutoffTime)
|
||
ct := time.Date(t.Year(), t.Month(), t.Day(), cutoff.Hour(), cutoff.Minute(), 0, 0, t.Location())
|
||
if t.Before(ct) {
|
||
return t.Format("2006-01-02")
|
||
}
|
||
return t.Add(24 * time.Hour).Format("2006-01-02")
|
||
}
|
||
|
||
func (s *Store) loadState() *batchState {
|
||
data, err := os.ReadFile(s.stateFile)
|
||
if err != nil {
|
||
return nil
|
||
}
|
||
var st batchState
|
||
if err := json.Unmarshal(data, &st); err != nil {
|
||
s.log(fmt.Sprintf("Failed to load state file: %v", err))
|
||
return nil
|
||
}
|
||
|
||
currentDate := s.countingDate(time.Now())
|
||
if st.CountingDate != currentDate {
|
||
s.log(fmt.Sprintf("State file belongs to previous counting day (%s). Finalizing before fresh start.", st.CountingDate))
|
||
s.insertBatch(st.CountingDate, st.BatchNumber, st.Count, st.StartTime, time.Now().Format(time.RFC3339))
|
||
os.Remove(s.stateFile)
|
||
return nil
|
||
}
|
||
|
||
s.log(fmt.Sprintf("Resumed batch #%d from %s with count=%d", st.BatchNumber, st.StartTime, st.Count))
|
||
s.resetBatchTimer()
|
||
return &st
|
||
}
|
||
|
||
func (s *Store) saveState() {
|
||
if s.currentState == nil {
|
||
os.Remove(s.stateFile)
|
||
return
|
||
}
|
||
data, err := json.MarshalIndent(s.currentState, "", " ")
|
||
if err != nil {
|
||
return
|
||
}
|
||
if err := os.WriteFile(s.stateFile, data, 0644); err != nil {
|
||
s.log(fmt.Sprintf("Failed to save state: %v", err))
|
||
}
|
||
}
|
||
|
||
func (s *Store) nextBatchNumber(countingDate string) int {
|
||
var maxNum sql.NullInt64
|
||
err := s.db.QueryRow(
|
||
`SELECT MAX(batch_number) FROM batches WHERE counting_date = ? AND camera_name = ? AND object_label = ?`,
|
||
countingDate, s.cameraName, s.objectLabel,
|
||
).Scan(&maxNum)
|
||
if err != nil || !maxNum.Valid {
|
||
return 1
|
||
}
|
||
return int(maxNum.Int64) + 1
|
||
}
|
||
|
||
func (s *Store) startNewBatch(countingDate string) {
|
||
batchNum := s.nextBatchNumber(countingDate)
|
||
now := time.Now().Format(time.RFC3339)
|
||
|
||
var countedIDs []string
|
||
if s.previousState != nil && len(s.previousState.CountedEventIDs) > 0 {
|
||
ids := s.previousState.CountedEventIDs
|
||
start := 0
|
||
if len(ids) > s.carryIDs {
|
||
start = len(ids) - s.carryIDs
|
||
}
|
||
countedIDs = ids[start:]
|
||
}
|
||
|
||
s.currentState = &batchState{
|
||
CountingDate: countingDate,
|
||
BatchNumber: batchNum,
|
||
Count: 0,
|
||
StartTime: now,
|
||
LastDetection: now,
|
||
CountedEventIDs: countedIDs,
|
||
}
|
||
s.saveState()
|
||
s.log(fmt.Sprintf("Started batch #%d for %s (%s)", batchNum, countingDate, s.objectLabel))
|
||
}
|
||
|
||
func (s *Store) resetBatchTimer() {
|
||
if s.batchTimer != nil {
|
||
s.batchTimer.Stop()
|
||
}
|
||
s.batchTimer = time.AfterFunc(s.batchTimeout, func() {
|
||
s.log(fmt.Sprintf("Batch inactivity timeout (%.0fs) reached", s.batchTimeout.Seconds()))
|
||
s.endBatch("timeout")
|
||
})
|
||
}
|
||
|
||
func (s *Store) startIgnoreBatchLabel() {
|
||
if s.ignoreBatchLabelTimer != nil {
|
||
return
|
||
}
|
||
s.ignoreBatchLabel = true
|
||
s.ignoreBatchLabelTimer = time.AfterFunc(s.ignoreBatchLabelTimeout, func() {
|
||
s.ignoreBatchLabel = false
|
||
s.ignoreBatchLabelTimer = nil
|
||
s.log("Ignore batch label cooldown finished")
|
||
})
|
||
s.log(fmt.Sprintf("Ignore batch label for %.0fs", s.ignoreBatchLabelTimeout.Seconds()))
|
||
}
|
||
|
||
func (s *Store) RecordAyamCrossing(trackID int) (int, bool) {
|
||
s.mu.Lock()
|
||
defer s.mu.Unlock()
|
||
|
||
countingDate := s.countingDate(time.Now())
|
||
startedNew := false
|
||
|
||
if s.currentState == nil {
|
||
s.startNewBatch(countingDate)
|
||
startedNew = true
|
||
} else if s.currentState.CountingDate != countingDate {
|
||
s.endBatchLocked("cutoff")
|
||
s.startNewBatch(countingDate)
|
||
startedNew = true
|
||
}
|
||
|
||
eventKey := fmt.Sprintf("%d", trackID)
|
||
found := false
|
||
for _, id := range s.currentState.CountedEventIDs {
|
||
if id == eventKey {
|
||
found = true
|
||
break
|
||
}
|
||
}
|
||
if !found {
|
||
s.currentState.Count++
|
||
s.currentState.CountedEventIDs = append(s.currentState.CountedEventIDs, eventKey)
|
||
s.log(fmt.Sprintf("Counted ayam (track %d) | batch #%d total: %d", trackID, s.currentState.BatchNumber, s.currentState.Count))
|
||
}
|
||
|
||
s.currentState.LastDetection = time.Now().Format(time.RFC3339)
|
||
s.saveState()
|
||
s.resetBatchTimer()
|
||
return s.currentState.Count, startedNew
|
||
}
|
||
|
||
func (s *Store) RecordTalenanCrossing(trackID int) bool {
|
||
if s.ignoreBatchLabel {
|
||
return false
|
||
}
|
||
s.mu.Lock()
|
||
defer s.mu.Unlock()
|
||
|
||
s.startIgnoreBatchLabel()
|
||
s.endBatchLocked("talenan")
|
||
s.log(fmt.Sprintf("Batch closed by talenan (track %d)", trackID))
|
||
|
||
if s.batchTimer != nil {
|
||
s.batchTimer.Stop()
|
||
s.batchTimer = nil
|
||
}
|
||
return true
|
||
}
|
||
|
||
func (s *Store) endBatch(closedBy string) {
|
||
s.mu.Lock()
|
||
defer s.mu.Unlock()
|
||
s.endBatchLocked(closedBy)
|
||
}
|
||
|
||
func (s *Store) endBatchLocked(closedBy string) bool {
|
||
if s.currentState == nil {
|
||
return false
|
||
}
|
||
|
||
s.previousState = s.currentState
|
||
state := s.currentState
|
||
|
||
startTime, _ := time.Parse(time.RFC3339, state.StartTime)
|
||
endTime := time.Now()
|
||
durationSec := endTime.Sub(startTime).Seconds()
|
||
|
||
if state.Count < s.minObjectPerBatch || int(durationSec) < s.minDurationPerBatch {
|
||
s.currentState = nil
|
||
s.saveState()
|
||
if s.batchTimer != nil {
|
||
s.batchTimer.Stop()
|
||
s.batchTimer = nil
|
||
}
|
||
s.log(fmt.Sprintf("Batch #%d discarded (count=%d, duration=%.0fs)", state.BatchNumber, state.Count, durationSec))
|
||
return false
|
||
}
|
||
|
||
endTimeStr := endTime.Format(time.RFC3339)
|
||
if err := s.insertBatch(state.CountingDate, state.BatchNumber, state.Count, state.StartTime, endTimeStr); err != nil {
|
||
s.log(fmt.Sprintf("Failed to persist batch: %v", err))
|
||
return false
|
||
}
|
||
|
||
cps := float64(state.Count) / durationSec
|
||
if durationSec == 0 {
|
||
cps = 0
|
||
}
|
||
s.log(fmt.Sprintf("Batch #%d ended | count=%d | duration=%.0fs | cps=%.3f | closed_by=%s",
|
||
state.BatchNumber, state.Count, durationSec, cps, closedBy))
|
||
|
||
s.currentState = nil
|
||
s.saveState()
|
||
if s.batchTimer != nil {
|
||
s.batchTimer.Stop()
|
||
s.batchTimer = nil
|
||
}
|
||
return true
|
||
}
|
||
|
||
func (s *Store) insertBatch(countingDate string, batchNumber, count int, startTime, endTime string) error {
|
||
tx, err := s.db.Begin()
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer tx.Rollback()
|
||
|
||
_, err = tx.Exec(
|
||
`INSERT INTO batches (counting_date, batch_number, camera_name, object_label, count, start_time, end_time)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?)`,
|
||
countingDate, batchNumber, s.cameraName, s.objectLabel, count, startTime, endTime,
|
||
)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
_, err = tx.Exec(
|
||
`INSERT INTO daily_summaries (counting_date, camera_name, object_label, total_count, total_batches)
|
||
VALUES (?, ?, ?, ?, 1)
|
||
ON CONFLICT(counting_date, camera_name, object_label)
|
||
DO UPDATE SET total_count = total_count + excluded.total_count,
|
||
total_batches = total_batches + excluded.total_batches,
|
||
updated_at = CURRENT_TIMESTAMP`,
|
||
countingDate, s.cameraName, s.objectLabel, count,
|
||
)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
if err := tx.Commit(); err != nil {
|
||
return err
|
||
}
|
||
|
||
var totalCount, totalBatches int
|
||
_ = s.db.QueryRow(
|
||
`SELECT total_count, total_batches FROM daily_summaries
|
||
WHERE counting_date = ? AND camera_name = ? AND object_label = ?`,
|
||
countingDate, s.cameraName, s.objectLabel,
|
||
).Scan(&totalCount, &totalBatches)
|
||
|
||
s.log(fmt.Sprintf("Daily totals for %s: %d objects across %d batch(es)", countingDate, totalCount, totalBatches))
|
||
return nil
|
||
}
|
||
|
||
func (s *Store) cutoffWatcher() {
|
||
ticker := time.NewTicker(60 * time.Second)
|
||
defer ticker.Stop()
|
||
for {
|
||
select {
|
||
case <-s.shutdownCh:
|
||
return
|
||
case <-ticker.C:
|
||
s.mu.Lock()
|
||
if s.currentState != nil && s.currentState.CountingDate != s.countingDate(time.Now()) {
|
||
s.log("Daily cutoff reached – finalizing batch")
|
||
s.endBatchLocked("cutoff")
|
||
}
|
||
s.mu.Unlock()
|
||
}
|
||
}
|
||
}
|
||
|
||
func (s *Store) StartCutoffWatcher() {
|
||
go s.cutoffWatcher()
|
||
}
|
||
|
||
func (s *Store) CurrentBatchNumber() int {
|
||
s.mu.Lock()
|
||
defer s.mu.Unlock()
|
||
if s.currentState == nil {
|
||
return 0
|
||
}
|
||
return s.currentState.BatchNumber
|
||
}
|
||
|
||
func (s *Store) CurrentBatchCount() int {
|
||
s.mu.Lock()
|
||
defer s.mu.Unlock()
|
||
if s.currentState == nil {
|
||
return 0
|
||
}
|
||
return s.currentState.Count
|
||
}
|
||
|
||
func (s *Store) ClosedTotalForDay() int {
|
||
countingDate := s.countingDate(time.Now())
|
||
var total int
|
||
_ = s.db.QueryRow(
|
||
`SELECT COALESCE(total_count, 0) FROM daily_summaries
|
||
WHERE counting_date = ? AND camera_name = ? AND object_label = ?`,
|
||
countingDate, s.cameraName, s.objectLabel,
|
||
).Scan(&total)
|
||
return total
|
||
}
|
||
|
||
func (s *Store) DisplayTotal() int {
|
||
return s.ClosedTotalForDay() + s.CurrentBatchCount()
|
||
}
|
||
|
||
func (s *Store) Shutdown() {
|
||
if s.stopped {
|
||
return
|
||
}
|
||
s.stopped = true
|
||
close(s.shutdownCh)
|
||
s.endBatch("shutdown")
|
||
if s.batchTimer != nil {
|
||
s.batchTimer.Stop()
|
||
}
|
||
if s.db != nil {
|
||
s.db.Close()
|
||
}
|
||
}
|