From: Matt Joiner <anacrolix@gmail.com>
Date: Tue, 3 Nov 2020 04:17:12 +0000 (+1100)
Subject: sqlite storage: Add capacity management
X-Git-Tag: v1.19.0~42
X-Git-Url: http://www.git.stargrave.org/?a=commitdiff_plain;h=96b640065a9f7d7e7e4e09d77b5a38faf973cff3;p=btrtrc.git

sqlite storage: Add capacity management
---

diff --git a/storage/sqlite/sqlite-storage.go b/storage/sqlite/sqlite-storage.go
index c0e59ccb..7eeb8a10 100644
--- a/storage/sqlite/sqlite-storage.go
+++ b/storage/sqlite/sqlite-storage.go
@@ -155,6 +155,8 @@ type NewPoolOpts struct {
 	// Forces WAL, disables shared caching.
 	ConcurrentBlobReads bool
 	DontInitSchema      bool
+	// If non-zero, overrides the existing setting.
+	Capacity int64
 }
 
 // There's some overlap here with NewPoolOpts, and I haven't decided what needs to be done. For now,
@@ -167,6 +169,16 @@ type ProviderOpts struct {
 	BatchWrites        bool
 }
 
+// Remove any capacity limits.
+func UnlimitCapacity(conn conn) error {
+	return sqlitex.Exec(conn, "delete from setting where key='capacity'", nil)
+}
+
+// Set the capacity limit to exactly this value.
+func SetCapacity(conn conn, cap int64) error {
+	return sqlitex.Exec(conn, "insert into setting values ('capacity', ?)", nil, cap)
+}
+
 func NewPool(opts NewPoolOpts) (_ ConnPool, _ ProviderOpts, err error) {
 	if opts.NumConns == 0 {
 		opts.NumConns = runtime.NumCPU()
@@ -191,14 +203,25 @@ func NewPool(opts NewPoolOpts) (_ ConnPool, _ ProviderOpts, err error) {
 	if err != nil {
 		return
 	}
+	defer func() {
+		if err != nil {
+			conns.Close()
+		}
+	}()
+	conn := conns.Get(context.TODO())
+	defer conns.Put(conn)
 	if !opts.DontInitSchema {
-		conn := conns.Get(context.TODO())
-		defer conns.Put(conn)
 		err = initSchema(conn)
 		if err != nil {
 			return
 		}
 	}
+	if opts.Capacity != 0 {
+		err = SetCapacity(conn, opts.Capacity)
+		if err != nil {
+			return
+		}
+	}
 	return conns, ProviderOpts{
 		NumConns:           opts.NumConns,
 		ConcurrentBlobRead: opts.ConcurrentBlobReads,