storage.go

package main

import (
	"bufio"
	"encoding/json"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"time"
)

var (
	dataDir    = filepath.Join(homeDir(), ".netwatch")
	usageFile  = filepath.Join(dataDir, "usage.jsonl")
	appsFile   = filepath.Join(dataDir, "apps.jsonl")
	spikesFile = filepath.Join(dataDir, "spikes.jsonl")
	lockFile   = filepath.Join(dataDir, "writer.lock")
)

const lockStale = 15
const appsPerMinute = 15

func homeDir() string { h, _ := os.UserHomeDir(); return h }

type rec struct {
	T    int64  `json:"t"`
	App  string `json:"app,omitempty"`
	Down int64  `json:"down"`
	Up   int64  `json:"up"`
}

type dayRow struct {
	Day        string
	Down, Up   int64
}

type appRow struct {
	App      string
	Down, Up int64
}

type UsageLog struct {
	bucket    int64
	bd, bu    int64
	appBucket int64
	appAcc    map[string][2]int64
	IsWriter  bool
	Session   [2]int64
	Today     [2]int64
	Week      [2]int64
}

func NewUsageLog() *UsageLog {
	os.MkdirAll(dataDir, 0755)
	l := &UsageLog{appAcc: map[string][2]int64{}}
	l.IsWriter = l.acquire()
	l.loadTotals()
	return l
}

func (l *UsageLog) acquire() bool {
	now := time.Now().Unix()
	if fi, err := os.Stat(lockFile); err == nil {
		if now-fi.ModTime().Unix() < lockStale {
			b, _ := os.ReadFile(lockFile)
			if strings.TrimSpace(string(b)) != strconv.Itoa(os.Getpid()) {
				return false
			}
		}
	}
	os.WriteFile(lockFile, []byte(strconv.Itoa(os.Getpid())), 0644)
	return true
}

func (l *UsageLog) refresh() {
	if l.IsWriter {
		os.WriteFile(lockFile, []byte(strconv.Itoa(os.Getpid())), 0644)
	}
}

func readRecs(path string) []rec {
	var out []rec
	f, err := os.Open(path)
	if err != nil {
		return out
	}
	defer f.Close()
	sc := bufio.NewScanner(f)
	sc.Buffer(make([]byte, 1024*1024), 1024*1024)
	for sc.Scan() {
		var r rec
		if json.Unmarshal(sc.Bytes(), &r) == nil {
			out = append(out, r)
		}
	}
	return out
}

func appendRec(path string, r rec) {
	f, err := os.OpenFile(path, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return
	}
	defer f.Close()
	b, _ := json.Marshal(r)
	f.Write(append(b, '\n'))
}

func startOfToday() int64 {
	n := time.Now()
	return time.Date(n.Year(), n.Month(), n.Day(), 0, 0, 0, 0, n.Location()).Unix()
}

func nowMinute() int64 { return time.Now().Unix() / 60 * 60 }

func (l *UsageLog) loadTotals() {
	st := startOfToday()
	wc := time.Now().Unix() - 7*86400
	for _, r := range readRecs(usageFile) {
		if r.T >= wc {
			l.Week[0] += r.Down
			l.Week[1] += r.Up
		}
		if r.T >= st {
			l.Today[0] += r.Down
			l.Today[1] += r.Up
		}
	}
}

func (l *UsageLog) Add(d, u int64) {
	l.Session[0] += d
	l.Session[1] += u
	l.Today[0] += d
	l.Today[1] += u
	l.Week[0] += d
	l.Week[1] += u
	m := nowMinute()
	if l.bucket == 0 {
		l.bucket = m
	}
	if m != l.bucket {
		l.flush()
		l.bucket = m
	}
	l.bd += d
	l.bu += u
	l.refresh()
}

func (l *UsageLog) flush() {
	if !l.IsWriter || l.bucket == 0 || (l.bd == 0 && l.bu == 0) {
		l.bd, l.bu = 0, 0
		return
	}
	appendRec(usageFile, rec{T: l.bucket, Down: l.bd, Up: l.bu})
	l.bd, l.bu = 0, 0
}

func (l *UsageLog) AddApps(ab map[string][2]int64) {
	m := nowMinute()
	if l.appBucket == 0 {
		l.appBucket = m
	}
	if m != l.appBucket {
		l.flushApps()
		l.appBucket = m
	}
	for n, v := range ab {
		c := l.appAcc[n]
		l.appAcc[n] = [2]int64{c[0] + v[0], c[1] + v[1]}
	}
}

func (l *UsageLog) flushApps() {
	if !l.IsWriter || l.appBucket == 0 || len(l.appAcc) == 0 {
		l.appAcc = map[string][2]int64{}
		return
	}
	type kv struct {
		n    string
		d, u int64
	}
	var arr []kv
	for n, v := range l.appAcc {
		arr = append(arr, kv{n, v[0], v[1]})
	}
	sort.Slice(arr, func(i, j int) bool { return arr[i].d+arr[i].u > arr[j].d+arr[j].u })
	if len(arr) > appsPerMinute {
		arr = arr[:appsPerMinute]
	}
	f, _ := os.OpenFile(appsFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if f != nil {
		for _, e := range arr {
			if e.d+e.u > 0 {
				b, _ := json.Marshal(rec{T: l.appBucket, App: e.n, Down: e.d, Up: e.u})
				f.Write(append(b, '\n'))
			}
		}
		f.Close()
	}
	l.appAcc = map[string][2]int64{}
}

func (l *UsageLog) LogSpike(app string, d, u float64) {
	if !l.IsWriter {
		return
	}
	appendRec(spikesFile, rec{T: time.Now().Unix(), App: app, Down: int64(d), Up: int64(u)})
}

func (l *UsageLog) Close() {
	l.flush()
	l.flushApps()
}

func DailySummary() []dayRow {
	m := map[string][2]int64{}
	for _, r := range readRecs(usageFile) {
		day := time.Unix(r.T, 0).Format("2006-01-02")
		c := m[day]
		m[day] = [2]int64{c[0] + r.Down, c[1] + r.Up}
	}
	var out []dayRow
	for d, v := range m {
		out = append(out, dayRow{d, v[0], v[1]})
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Day < out[j].Day })
	return out
}

func AppSummary(sinceSec int64) []appRow {
	cut := time.Now().Unix() - sinceSec
	m := map[string][2]int64{}
	for _, r := range readRecs(appsFile) {
		if r.T >= cut {
			c := m[r.App]
			m[r.App] = [2]int64{c[0] + r.Down, c[1] + r.Up}
		}
	}
	var out []appRow
	for a, v := range m {
		out = append(out, appRow{a, v[0], v[1]})
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Down+out[i].Up > out[j].Down+out[j].Up })
	return out
}

func RecentSpikes(sinceSec int64) []rec {
	cut := time.Now().Unix() - sinceSec
	var out []rec
	for _, r := range readRecs(spikesFile) {
		if r.T >= cut {
			out = append(out, r)
		}
	}
	return out
}

func CycleUsage(cycleDay int) (int64, int64) {
	n := time.Now()
	cd := cycleDay
	if cd < 1 {
		cd = 1
	}
	if cd > 28 {
		cd = 28
	}
	var start time.Time
	if n.Day() >= cd {
		start = time.Date(n.Year(), n.Month(), cd, 0, 0, 0, 0, n.Location())
	} else {
		pm := n.AddDate(0, -1, 0)
		start = time.Date(pm.Year(), pm.Month(), cd, 0, 0, 0, 0, pm.Location())
	}
	s := start.Unix()
	var d, u int64
	for _, r := range readRecs(usageFile) {
		if r.T >= s {
			d += r.Down
			u += r.Up
		}
	}
	return d, u
}