jarl/internal/stats/store.go

// Package stats collects query statistics from named: counters from its
// statistics channel and per-query data from its response log. It stores
// aggregates (never individual queries) in a SQLite database of its own,
// separate from Jarl's configuration database, so it can be deleted at any
// time without losing anything important.
package stats

import (
	"context"
	"database/sql"
	"errors"
	"fmt"

	_ "modernc.org/sqlite" // registers the "sqlite" driver
)

// Resolutions of stored data.
const (
	ResMinute = "m"
	ResHour   = "h"
	ResDay    = "d"
)

// Store is the statistics database.
type Store struct {
	db *sql.DB
}

var schema = []string{
	`CREATE TABLE counters (
		res   TEXT NOT NULL,
		ts    INTEGER NOT NULL,
		key   TEXT NOT NULL,
		value INTEGER NOT NULL,
		PRIMARY KEY (res, ts, key)
	) WITHOUT ROWID;
	CREATE INDEX counters_by_key ON counters (res, key, ts);
	CREATE TABLE tops (
		res   TEXT NOT NULL,
		ts    INTEGER NOT NULL,
		kind  TEXT NOT NULL,
		key   TEXT NOT NULL,
		count INTEGER NOT NULL,
		PRIMARY KEY (res, ts, kind, key)
	) WITHOUT ROWID;
	CREATE TABLE meta (key TEXT PRIMARY KEY, value TEXT NOT NULL);`,
}

// OpenStore opens (creating if needed) the statistics database. Use
// ":memory:" for tests.
func OpenStore(path string) (*Store, error) {
	dsn := "file:" + path + "?_pragma=busy_timeout(10000)&_pragma=synchronous(NORMAL)&_txlock=immediate"
	if path == ":memory:" {
		dsn = "file::memory:?_txlock=immediate"
	} else {
		dsn += "&_pragma=journal_mode(WAL)"
	}
	db, err := sql.Open("sqlite", dsn)
	if err != nil {
		return nil, err
	}
	if path == ":memory:" {
		db.SetMaxOpenConns(1)
	}
	s := &Store{db: db}
	if err := s.migrate(); err != nil {
		db.Close()
		return nil, fmt.Errorf("statistics database %s: %w", path, err)
	}
	return s, nil
}

// Close closes the database.
func (s *Store) Close() error { return s.db.Close() }

func (s *Store) migrate() error {
	if _, err := s.db.Exec(`CREATE TABLE IF NOT EXISTS schema_version (version INTEGER NOT NULL)`); err != nil {
		return err
	}
	var v int
	err := s.db.QueryRow(`SELECT version FROM schema_version`).Scan(&v)
	if errors.Is(err, sql.ErrNoRows) {
		if _, err := s.db.Exec(`INSERT INTO schema_version (version) VALUES (0)`); err != nil {
			return err
		}
	} else if err != nil {
		return err
	}
	if v > len(schema) {
		return fmt.Errorf("schema version %d is newer than this Jarl; upgrade Jarl or delete the statistics database", v)
	}
	for i := v; i < len(schema); i++ {
		tx, err := s.db.Begin()
		if err != nil {
			return err
		}
		if _, err := tx.Exec(schema[i]); err != nil {
			tx.Rollback()
			return err
		}
		if _, err := tx.Exec(`UPDATE schema_version SET version = ?`, i+1); err != nil {
			tx.Rollback()
			return err
		}
		if err := tx.Commit(); err != nil {
			return err
		}
	}
	return nil
}

// Batch collects writes that are committed together.
type Batch struct {
	counters []counterRow
	tops     []topsWrite
	meta     map[string]string
}

type counterRow struct {
	res   string
	ts    int64
	key   string
	value int64
}

type topsWrite struct {
	res, kind string
	ts        int64
	entries   []Entry
	replace   bool // replace the (res, ts, kind) rows instead of adding
	keep      int  // after adding, keep only the largest entries (0: all)
}

// Entry is a key and its count.
type Entry struct {
	Key   string `json:"key"`
	Count int64  `json:"count"`
	// Name is a client's host name, from Jarl's own records (set by the
	// API for client lists).
	Name string `json:"name,omitempty"`
}

// AddCounters adds values to the counters of a bucket.
func (b *Batch) AddCounters(res string, ts int64, values map[string]int64) {
	for k, v := range values {
		b.counters = append(b.counters, counterRow{res, ts, k, v})
	}
}

// SetTops replaces a bucket's top list of a kind.
func (b *Batch) SetTops(res string, ts int64, kind string, entries []Entry) {
	b.tops = append(b.tops, topsWrite{res: res, ts: ts, kind: kind, entries: entries, replace: true})
}

// AddTops adds counts to a bucket's top list and keeps the keep largest.
func (b *Batch) AddTops(res string, ts int64, kind string, entries []Entry, keep int) {
	b.tops = append(b.tops, topsWrite{res: res, ts: ts, kind: kind, entries: entries, keep: keep})
}

// SetMeta stores a value that survives restarts.
func (b *Batch) SetMeta(key, value string) {
	if b.meta == nil {
		b.meta = map[string]string{}
	}
	b.meta[key] = value
}

// Commit writes the batch in one transaction.
func (s *Store) Commit(ctx context.Context, b *Batch) error {
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback()
	if len(b.counters) > 0 {
		st, err := tx.PrepareContext(ctx, `INSERT INTO counters (res, ts, key, value) VALUES (?, ?, ?, ?)
			ON CONFLICT (res, ts, key) DO UPDATE SET value = value + excluded.value`)
		if err != nil {
			return err
		}
		for _, c := range b.counters {
			if _, err := st.ExecContext(ctx, c.res, c.ts, c.key, c.value); err != nil {
				st.Close()
				return err
			}
		}
		st.Close()
	}
	for _, t := range b.tops {
		if t.replace {
			if _, err := tx.ExecContext(ctx, `DELETE FROM tops WHERE res = ? AND ts = ? AND kind = ?`, t.res, t.ts, t.kind); err != nil {
				return err
			}
		}
		for _, e := range t.entries {
			if _, err := tx.ExecContext(ctx, `INSERT INTO tops (res, ts, kind, key, count) VALUES (?, ?, ?, ?, ?)
				ON CONFLICT (res, ts, kind, key) DO UPDATE SET count = count + excluded.count`, t.res, t.ts, t.kind, e.Key, e.Count); err != nil {
				return err
			}
		}
		if t.keep > 0 {
			if _, err := tx.ExecContext(ctx, `DELETE FROM tops WHERE res = ? AND ts = ? AND kind = ? AND key NOT IN (
				SELECT key FROM tops WHERE res = ? AND ts = ? AND kind = ? ORDER BY count DESC, key LIMIT ?)`,
				t.res, t.ts, t.kind, t.res, t.ts, t.kind, t.keep); err != nil {
				return err
			}
		}
	}
	for k, v := range b.meta {
		if _, err := tx.ExecContext(ctx, `INSERT INTO meta (key, value) VALUES (?, ?) ON CONFLICT (key) DO UPDATE SET value = excluded.value`, k, v); err != nil {
			return err
		}
	}
	return tx.Commit()
}

// Meta returns a stored value ("" if unset).
func (s *Store) Meta(ctx context.Context, key string) (string, error) {
	var v string
	err := s.db.QueryRowContext(ctx, `SELECT value FROM meta WHERE key = ?`, key).Scan(&v)
	if errors.Is(err, sql.ErrNoRows) {
		return "", nil
	}
	return v, err
}

// Prune deletes data of a resolution older than before (a Unix time).
func (s *Store) Prune(ctx context.Context, res string, before int64) (int64, error) {
	var n int64
	for _, table := range []string{"counters", "tops"} {
		r, err := s.db.ExecContext(ctx, `DELETE FROM `+table+` WHERE res = ? AND ts < ?`, res, before)
		if err != nil {
			return n, err
		}
		c, _ := r.RowsAffected()
		n += c
	}
	return n, nil
}

// TopsAt returns a bucket's top list of a kind.
func (s *Store) TopsAt(ctx context.Context, res string, ts int64, kind string) ([]Entry, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT key, count FROM tops WHERE res = ? AND ts = ? AND kind = ? ORDER BY count DESC, key`, res, ts, kind)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []Entry
	for rows.Next() {
		var e Entry
		if err := rows.Scan(&e.Key, &e.Count); err != nil {
			return nil, err
		}
		out = append(out, e)
	}
	return out, rows.Err()
}