misc/infrastructure/powerbot/db.go

   1 package main
   2
   3 import (
   4         "database/sql"
   5         "fmt"
   6         "maunium.net/go/mautrix/id"
   7         _ "modernc.org/sqlite"
   8         "time"
   9 )
  10
  11 type State int
  12
  13 const (
  14         NotActive State = iota
  15         Idle
  16         Active
  17 )
  18
  19 type Score struct {
  20         LastEvent    time.Time
  21         CurrentState State
  22         Idle         time.Duration
  23         Active       time.Duration
  24 }
  25
  26 const dbSchema = `
  27 CREATE TABLE IF NOT EXISTS room_users (
  28   room_id STRING NOT NULL,
  29   user_id STRING NOT NULL,
  30   state_time TIMESTAMP NOT NULL,
  31   state INT NOT NULL,
  32   idle_nsec INT64 NOT NULL,
  33   active_nsec INT64 NOT NULL,
  34   PRIMARY KEY(room_id, user_id)
  35 );
  36 `
  37
  38 const fetchStateQuery = `
  39 SELECT state_time, state, idle_nsec, active_nsec
  40 FROM room_users
  41 WHERE room_id = ?
  42   AND user_id = ?
  43 `
  44
  45 const insertStateQuery = `
  46 INSERT INTO room_users(room_id, user_id, state_time, state, idle_nsec, active_nsec)
  47 VALUES(?, ?, ?, ?, 0.0, 0.0)
  48 `
  49
  50 const updateStateQuery = `
  51 UPDATE room_users
  52 SET state_time = ?, state = ?, idle_nsec = ?, active_nsec = ?
  53 WHERE room_id = ?
  54   AND user_id = ?
  55 `
  56
  57 const fetchUserScoresQuery = `
  58 SELECT user_id, state_time, state, idle_nsec, active_nsec
  59 FROM room_users
  60 WHERE room_id = ?
  61 `
  62
  63 var db *sql.DB
  64
  65 func InitDatabase() error {
  66         var err error
  67         db, err = sql.Open("sqlite", "users.sqlite")
  68         if err != nil {
  69                 return fmt.Errorf("could not open SQLite database: %v", err)
  70         }
  71         _, err = db.Exec(dbSchema)
  72         if err != nil {
  73                 return fmt.Errorf("could not set SQLite database schema: %v", err)
  74         }
  75         return nil
  76 }
  77
  78 func CloseDatabase() error {
  79         return db.Close()
  80 }
  81
  82 func queryUserScores(room id.RoomID, now time.Time) (map[id.UserID]*Score, error) {
  83         var users map[id.UserID]*Score
  84         err := retryPolicy(func() error {
  85                 rows, err := db.Query(fetchUserScoresQuery, room)
  86                 if err != nil {
  87                         return fmt.Errorf("could not query users: %v", err)
  88                 }
  89                 users = map[id.UserID]*Score{}
  90                 for rows.Next() {
  91                         var user id.UserID
  92                         var score Score
  93                         if err := rows.Scan(&user, &score.LastEvent, &score.CurrentState, &score.Idle, &score.Active); err != nil {
  94                                 return fmt.Errorf("could not scan users query result: %v", err)
  95                         }
  96                         newScore := advanceScore(score, now)
  97                         users[user] = &newScore
  98                 }
  99                 if err := rows.Err(); err != nil {
 100                         return fmt.Errorf("could not read users: %v", err)
 101                 }
 102                 return nil
 103         })
 104         return users, err
 105 }
 106
 107 func advanceScore(score Score, now time.Time) Score {
 108         if !now.After(score.LastEvent) {
 109                 return score
 110         }
 111         dt := now.Sub(score.LastEvent)
 112         switch score.CurrentState {
 113         case Idle:
 114                 score.Idle += dt
 115         case Active:
 116                 score.Active += dt
 117         }
 118         return score
 119 }
 120
 121 func retryPolicy(f func() error) error {
 122         var err error
 123         for attempt := 0; attempt < 12; attempt++ {
 124                 err = f()
 125                 if err == nil {
 126                         return nil
 127                 }
 128                 time.Sleep(time.Millisecond * time.Duration(1<<attempt))
 129         }
 130         return err
 131 }
 132
 133 func inTx(db *sql.DB, f func(tx *sql.Tx) error) error {
 134         tx, err := db.Begin()
 135         if err != nil {
 136                 return fmt.Errorf("failed to create transaction: %v", err)
 137         }
 138         err = f(tx)
 139         if err != nil {
 140                 tx.Rollback()
 141                 return err
 142         }
 143         return tx.Commit()
 144 }
 145
 146 func writeUserStateAt(room id.RoomID, user id.UserID, now time.Time, maxPrevState, state State) error {
 147         return retryPolicy(func() error {
 148                 return inTx(db, func(tx *sql.Tx) error {
 149                         row := tx.QueryRow(fetchStateQuery, room, user)
 150                         var score Score
 151                         err := row.Scan(&score.LastEvent, &score.CurrentState, &score.Idle, &score.Active)
 152                         if err == sql.ErrNoRows {
 153                                 _, err = tx.Exec(insertStateQuery, room, user, now, state)
 154                                 if err != nil {
 155                                         return fmt.Errorf("failed to set state for new user: %v", err)
 156                                 }
 157                                 return nil
 158                         } else {
 159                                 if err != nil {
 160                                         return fmt.Errorf("failed to fetch state for user: %v", err)
 161                                 }
 162                                 if now.After(score.LastEvent) {
 163                                         if score.CurrentState > maxPrevState {
 164                                                 score.CurrentState = maxPrevState
 165                                         }
 166                                         score = advanceScore(score, now)
 167                                         _, err = tx.Exec(updateStateQuery, now, state, score.Idle, score.Active, room, user)
 168                                         if err != nil {
 169                                                 return fmt.Errorf("failed to update state for new user: %v", err)
 170                                         }
 171                                         return nil
 172                                 } else {
 173                                         _, err = tx.Exec(updateStateQuery, score.LastEvent, state, score.Idle, score.Active, room, user)
 174                                         if err != nil {
 175                                                 return fmt.Errorf("failed to update state for new user: %v", err)
 176                                         }
 177                                         return nil
 178                                 }
 179                         }
 180                 })
 181         })
 182 }