1
package sqlite3
2
 
3
import (
4
	"context"
5
	"database/sql"
6
	"time"
7
 
8
	"go.miragespace.co/specter/spec/chord"
9
)
10
 
11
func durationGuard(t time.Duration) (time.Duration, bool) {
12
	td := t.Truncate(time.Second)
13
	if td < time.Second {
14
		return 0, false
15
	}
16
	return td, true
17
}
18
 
19
func (s *SqliteKV) Acquire(ctx context.Context, lease []byte, ttl time.Duration) (uint64, error) {
20
	ttl, ok := durationGuard(ttl)
21
	if !ok {
22
		return 0, chord.ErrKVLeaseInvalidTTL
23
	}
24
 
25
	var next uint64
26
	err := withWriteTx(ctx, s.writer, func(tx *sql.Tx) error {
27
		now := time.Now()
28
		next = uint64(now.Add(ttl).UnixNano())
29
 
30
		res, err := tx.StmtContext(ctx, s.stmts.leaseAcquire).Exec(lease, bindUint64AsInt64(next), bindUint64AsInt64(uint64(now.UnixNano())))
31
		if err != nil {
32
			return err
33
		}
34
		n, err := res.RowsAffected()
35
		if err != nil {
36
			return err
37
		}
38
		if n == 0 {
39
			return chord.ErrKVLeaseConflict
40
		}
41
		return s.updateKeyTracker(ctx, tx, lease, LeaseFlag, 0)
42
	})
43
	if err != nil {
44
		return 0, err
45
	}
46
	return next, nil
47
}
48
 
49
func (s *SqliteKV) Renew(ctx context.Context, lease []byte, ttl time.Duration, prevToken uint64) (uint64, error) {
50
	ttl, ok := durationGuard(ttl)
51
	if !ok {
52
		return 0, chord.ErrKVLeaseInvalidTTL
53
	}
54
 
55
	var next uint64
56
	err := withWriteTx(ctx, s.writer, func(tx *sql.Tx) error {
57
		now := time.Now()
58
		next = uint64(time.Now().Add(ttl).UnixNano())
59
 
60
		res, err := tx.StmtContext(ctx, s.stmts.leaseRenew).Exec(
61
			bindUint64AsInt64(next),
62
			lease,
63
			bindUint64AsInt64(prevToken),
64
			bindUint64AsInt64(uint64(now.UnixNano())),
65
		)
66
		if err != nil {
67
			return err
68
		}
69
		n, err := res.RowsAffected()
70
		if err != nil {
71
			return err
72
		}
73
		if n == 0 {
74
			return chord.ErrKVLeaseExpired
75
		}
76
		return s.updateKeyTracker(ctx, tx, lease, LeaseFlag, 0)
77
	})
78
	if err != nil {
79
		return 0, err
80
	}
81
	// Post-commit check: re-read the row in a new session
82
	//    If the database row's token is no longer what we set, it means
83
	//    some other transaction updated it after we did our compare-and-set.
84
	var curToken int64
85
	readErr := s.stmts.leaseGet.QueryRowContext(ctx, lease).Scan(&curToken)
86
	if readErr != nil {
87
		// If the row no longer exists or some other error,
88
		// you could consider that as concurrency override or just return readErr.
89
		return 0, readErr
90
	}
91
 
92
	// 3. Compare the token we set (`next`) vs. what's actually in the DB now
93
	if scanInt64AsUint64(curToken) != next {
94
		// Another transaction must have updated it post-commit => concurrency lost
95
		return 0, chord.ErrKVLeaseExpired
96
	}
97
	return next, nil
98
}
99
 
100
func (s *SqliteKV) Release(ctx context.Context, lease []byte, token uint64) error {
101
	return withWriteTx(ctx, s.writer, func(tx *sql.Tx) error {
102
		res, err := tx.StmtContext(ctx, s.stmts.leaseRelease).Exec(lease, bindUint64AsInt64(token))
103
		if err != nil {
104
			return err
105
		}
106
		n, err := res.RowsAffected()
107
		if err != nil {
108
			return err
109
		}
110
		if n == 0 {
111
			return chord.ErrKVLeaseExpired
112
		}
113
		return s.updateKeyTracker(ctx, tx, lease, 0, LeaseFlag)
114
	})
115
}