From 3ced08127192f81002c49350684f97b36efc6695 Mon Sep 17 00:00:00 2001
From: Seebs <seebs@molecula.com>
Date: Fri, 18 Feb 2022 11:32:31 -0600
Subject: [PATCH] prevent crashes when closing db

When closing, we need to wait for existing Tx to exit before truncating
files and unmapping things. This shouldn't matter, because we don't actually
close the DB until all transactions are done, normally... except for the
background usage-gathering task. But really, it's probably just better to
be conservative.

The actual logic is fancier than it looks. We can't hold db.mu.Lock during
this, or the existing Tx can't exit. So we first grab the lock, set the closed
flag, set up a waiter for all current Tx to exit, and then release the lock.
Now we wait on the current Tx exiting. Once that's done, we grab the locks.
Anything coming in that tries to start a Tx will fail out fairly quickly
because the opened flag is now false, so even if other things get those
locks before we do, they won't keep them or create new Tx.

This makes one test deadlock because it opens a Tx and never closes it,
so we change that test to close its Tx.
---
 fragment_internal_test.go |  1 +
 rbf/db.go                 | 18 +++++++--
 rbf/db_test.go            | 84 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 100 insertions(+), 3 deletions(-)

diff --git a/fragment_internal_test.go b/fragment_internal_test.go
index 6fd803954..3beea99ec 100644
--- a/fragment_internal_test.go
+++ b/fragment_internal_test.go
@@ -1570,6 +1570,7 @@ func TestFragment_WriteTo_ReadFrom(t *testing.T) {
 	}
 	// make a read-only Tx after ReadFrom has committed.
 	tx = idx.holder.txf.NewTx(Txo{Write: !writable, Index: idx, Fragment: f1, Shard: f1.shard})
+	defer tx.Rollback()
 
 	// Verify cache is in other fragment.
 	if n := f1.cache.Len(); n != 1 {
diff --git a/rbf/db.go b/rbf/db.go
index ec75a6064..bd92f2271 100644
--- a/rbf/db.go
+++ b/rbf/db.go
@@ -411,16 +411,28 @@ func (db *DB) checkpoint() (err error) {
 
 // Close closes the database.
 func (db *DB) Close() (err error) {
-	// TODO(bbj): Add wait group to hang until last Tx is complete.
+	// mark db as closed, spawn a thing to wait for existing tx to drain, then
+	// release the lock so they CAN drain. We do this before getting the
+	// write lock, so if something else is waiting on rwmu.Lock, and will be
+	// competing with us, we can ensure that it'll exit out quickly.
+	db.mu.Lock()
+	db.opened = false
+	// wait for transactions to complete
+	ch := make(chan struct{})
+	db.afterCurrentTx(func() {
+		close(ch)
+	})
+	db.mu.Unlock()
+	<-ch
+
 	// Wait for writer lock.
 	db.rwmu.Lock()
 	defer db.rwmu.Unlock()
 
+	// and main DB lock.
 	db.mu.Lock()
 	defer db.mu.Unlock()
 
-	db.opened = false
-
 	// Close mmap handle.
 	if db.data != nil {
 		if e := syswrap.Munmap(db.data); e != nil && err == nil {
diff --git a/rbf/db_test.go b/rbf/db_test.go
index 3e6bdd1b1..2b886677c 100644
--- a/rbf/db_test.go
+++ b/rbf/db_test.go
@@ -139,6 +139,90 @@ func TestDB_WAL(t *testing.T) {
 			t.Fatal(err)
 		}
 	})
+
+	// initially this is just a cut and paste of the Halt test, except that
+	// we close the DB while the reads are still running.
+	t.Run("Close", func(t *testing.T) {
+		if testing.Short() {
+			t.Skip("-short enabled, skipping")
+		}
+
+		config := rbfcfg.NewDefaultConfig()
+		config.MaxWALSize = 16 * rbf.PageSize
+		config.MaxWALCheckpointSize = 8 * rbf.PageSize
+		config.MinWALCheckpointSize = 4 * rbf.PageSize
+
+		db := MustOpenDB(t, config)
+
+		// Continuously run read overlapping transactions.
+		ctx, cancel := context.WithCancel(context.Background())
+		g, ctx := errgroup.WithContext(ctx)
+		for i := 0; i < 10; i++ {
+			i := i
+			g.Go(func() error {
+				time.Sleep(time.Duration(i) * 10 * time.Millisecond) // stagger
+				for {
+					if err := ctx.Err(); err != nil {
+						return nil
+					}
+
+					if err := func() error {
+						tx, err := db.Begin(false)
+						if err != nil {
+							return err
+						}
+						// give the db time to close between when we opened and
+						// when we run the Container call
+						time.Sleep(10 * time.Millisecond)
+						_, err = tx.Container("x", 0)
+						if err != nil {
+							t.Fatalf("requesting container: %v", err)
+						}
+						defer tx.Rollback()
+						return nil
+					}(); err != nil {
+						// it's okay to ErrClosed, because we plan to close
+						// the database out from under us.
+						if err != rbf.ErrClosed {
+							return err
+						} else {
+							return nil
+						}
+					}
+				}
+			})
+		}
+
+		// Generate updates to the DB/WAL.
+		for i := 0; i < 100; i++ {
+			func() {
+				tx := MustBegin(t, db, true)
+				defer tx.Rollback()
+
+				if err := tx.CreateBitmapIfNotExists("x"); err != nil {
+					t.Fatal(err)
+				} else if _, err := tx.Add("x", uint64(i)); err != nil {
+					t.Fatal(err)
+				} else if err := tx.Commit(); err != nil {
+					t.Fatal(err)
+				}
+				time.Sleep(1 * time.Millisecond)
+			}()
+		}
+		// close the db now.
+		err := db.Close()
+		if err != nil {
+			t.Fatalf("closing db: %v", err)
+		}
+		// delay a bit to let some readers try to read
+		time.Sleep(20 * time.Millisecond)
+
+		// Stop read transactions & wait.
+		cancel()
+		if err := g.Wait(); err != nil {
+			t.Fatal(err)
+		}
+	})
 }
 
 func TestDB_Recovery(t *testing.T) {