First commit
This commit is contained in:
proitlab committed
1 parent
4f70c7f8d5
commit
dc8dcca75e
17 files changed
+4382
-2
No files matched your search
@@ -0,0 +1,467 @@
|
||||
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()
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user