CBG-3271 make sure expiration works on a closed bucket

bucket.go

@@ -36,13 +36,12 @@ type Bucket struct {
 	name            string         // Bucket name
 	collections     collectionsMap // Collections, indexed by DataStoreName
 	collectionFeeds map[sgbucket.DataStoreNameImpl][]*dcpFeed
-	mutex           *sync.Mutex // mutex for synchronized access to Bucket
-	sqliteDB        *sql.DB     // SQLite database handle (do not access; call db() instead)
-	expTimer        *time.Timer // Schedules expiration of docs
-	nextExp         *uint32     // Timestamp when expTimer will run (0 if never)
-	serial          uint32      // Serial number for logging
-	inMemory        bool        // True if it's an in-memory database
-	closed          bool        // represents state when it is closed
+	mutex           *sync.Mutex    // mutex for synchronized access to Bucket
+	sqliteDB        *sql.DB        // SQLite database handle (do not access; call db() instead)
+	expManager      *expiryManager // expiration manager for bucket
+	serial          uint32         // Serial number for logging
+	inMemory        bool           // True if it's an in-memory database
+	closed          bool           // represents state when it is closed
 }
 
 type collectionsMap = map[sgbucket.DataStoreNameImpl]*Collection
@@ -87,6 +86,7 @@ func OpenBucket(urlStr string, bucketName string, mode OpenMode) (b *Bucket, err
 	traceEnter("OpenBucket", "%q, %d", urlStr, mode)
 	defer func() { traceExit("OpenBucket", err, "ok") }()
 
+	ctx := context.TODO()
 	u, err := encodeDBURL(urlStr)
 	if err != nil {
 		return nil, err
@@ -95,13 +95,18 @@ func OpenBucket(urlStr string, bucketName string, mode OpenMode) (b *Bucket, err
 
 	bucket := getCachedBucket(bucketName)
 	if bucket != nil {
+		defer func() {
+			if err != nil {
+				bucket.Close(ctx)
+			}
+		}()
+
 		if mode == CreateNew {
 			return nil, fs.ErrExist
 		}
 		if urlStr != bucket.url {
 			return nil, fmt.Errorf("bucket %q already exists at %q, will not open at %q", bucketName, bucket.url, urlStr)
 		}
-		registerBucket(bucket)
 		return bucket, nil
 
 	}
@@ -174,13 +179,13 @@ func OpenBucket(urlStr string, bucketName string, mode OpenMode) (b *Bucket, err
 		collections:     make(map[sgbucket.DataStoreNameImpl]*Collection),
 		collectionFeeds: make(map[sgbucket.DataStoreNameImpl][]*dcpFeed),
 		mutex:           &sync.Mutex{},
-		nextExp:         new(uint32),
 		inMemory:        inMemory,
 		serial:          serial,
 	}
+	bucket.expManager = newExpirationManager(bucket.doExpiration)
 	defer func() {
 		if err != nil {
-			_ = bucket.CloseAndDelete(context.TODO())
+			_ = bucket.CloseAndDelete(ctx)
 		}
 	}()
 
@@ -194,16 +199,23 @@ func OpenBucket(urlStr string, bucketName string, mode OpenMode) (b *Bucket, err
 		if err = bucket.initializeSchema(bucketName); err != nil {
 			return nil, err
 		}
-	} else {
-		bucket.scheduleExpiration()
 	}
 	err = bucket.setName(bucketName)
 	if err != nil {
 		return nil, err
 	}
 
-	registerBucket(bucket)
-	return bucket, err
+	exists, bucketCopy := registerBucket(bucket)
+	// someone else beat registered the bucket in the registry, that's OK we'll close ours
+	if exists {
+		bucket.Close(ctx)
+	}
+	// only schedule expiration if bucket is not new. This doesn't need to be locked because only one bucket will execute this code.
+	if vers != 0 {
+		bucket._scheduleExpiration()
+	}
+
+	return bucketCopy, err
 }
 
 // Creates or re-opens a bucket, like OpenBucket.
@@ -307,7 +319,7 @@ func (bucket *Bucket) initializeSchema(bucketName string) (err error) {
 	return
 }
 
-// Returns the database handle as a `queryable` interface value.
+// db returns the database handle as a `queryable` interface value.
 // If the bucket has been closed, it returns a special `closedDB` value that will return
 // ErrBucketClosed from any call.
 func (bucket *Bucket) db() queryable {
@@ -316,7 +328,7 @@ func (bucket *Bucket) db() queryable {
 	return bucket._db()
 }
 
-// Returns the database handle as a `queryable` interface value. This is the same as `db()` without locking. This is not safe to call without bucket.mutex being locked the caller.
+// db returns the database handle as a `queryable` interface value. This is the same as `db()` without locking. This is not safe to call without bucket.mutex being locked the caller.
 func (bucket *Bucket) _db() queryable {
 	if bucket.closed {
 		return closedDB{}
@@ -380,8 +392,7 @@ func (b *Bucket) copy() *Bucket {
 		collections:     make(collectionsMap),
 		mutex:           b.mutex,
 		sqliteDB:        b.sqliteDB,
-		expTimer:        b.expTimer,
-		nextExp:         b.nextExp,
+		expManager:      b.expManager,
 		serial:          b.serial,
 		inMemory:        b.inMemory,
 	}

bucket_api.go

@@ -13,7 +13,6 @@ import (
 	"database/sql"
 	"errors"
 	"fmt"
-	"time"
 
 	sgbucket "github.com/couchbase/sg-bucket"
 )
@@ -62,9 +61,7 @@ func (bucket *Bucket) Close(_ context.Context) {
 
 // _closeSqliteDB closes the underlying sqlite database and shuts down dcpFeeds. Must have a lock to call this function.
 func (bucket *Bucket) _closeSqliteDB() {
-	if bucket.expTimer != nil {
-		bucket.expTimer.Stop()
-	}
+	bucket.expManager.stop()
 	for _, c := range bucket.collections {
 		c.close()
 	}
@@ -174,6 +171,7 @@ func (bucket *Bucket) DropDataStore(name sgbucket.DataStoreName) error {
 	return bucket.dropCollection(sc)
 }
 
+// ListDataStores returns a list of the names of all data stores in the bucket.
 func (bucket *Bucket) ListDataStores() (result []sgbucket.DataStoreName, err error) {
 	traceEnter("ListDataStores", "%s", bucket)
 	defer func() { traceExit("ListDataStores", err, "%v", result) }()
@@ -309,8 +307,8 @@ func (bucket *Bucket) dropCollection(name sgbucket.DataStoreNameImpl) error {
 
 //////// EXPIRATION (CUSTOM API):
 
-// Returns the earliest expiration time of any document, or 0 if none.
-func (bucket *Bucket) NextExpiration() (exp Exp, err error) {
+// nextExpiration returns the earliest expiration time of any document, or 0 if none.
+func (bucket *Bucket) nextExpiration() (exp Exp, err error) {
 	var expVal sql.NullInt64
 	row := bucket.db().QueryRow(`SELECT min(exp) FROM documents WHERE exp > 0`)
 	err = scan(row, &expVal)
@@ -320,8 +318,8 @@ func (bucket *Bucket) NextExpiration() (exp Exp, err error) {
 	return
 }
 
-// Immediately deletes all expired documents in this bucket.
-func (bucket *Bucket) ExpireDocuments() (int64, error) {
+// expireDocuments immediately deletes all expired documents in this bucket.
+func (bucket *Bucket) expireDocuments() (int64, error) {
 	names, err := bucket.ListDataStores()
 	if err != nil {
 		return 0, err
@@ -330,7 +328,7 @@ func (bucket *Bucket) ExpireDocuments() (int64, error) {
 	for _, name := range names {
 		if coll, err := bucket.getCollection(name.(sgbucket.DataStoreNameImpl)); err != nil {
 			return 0, err
-		} else if n, err := coll.ExpireDocuments(); err != nil {
+		} else if n, err := coll.expireDocuments(); err != nil {
 			return 0, err
 		} else {
 			count += n
@@ -339,45 +337,25 @@ func (bucket *Bucket) ExpireDocuments() (int64, error) {
 	return count, nil
 }
 
-func (bucket *Bucket) scheduleExpiration() {
-	if nextExp, err := bucket.NextExpiration(); err == nil && nextExp > 0 {
-		bucket.scheduleExpirationAtOrBefore(nextExp)
-	}
-}
-
-func (bucket *Bucket) scheduleExpirationAtOrBefore(exp uint32) {
-	if exp > 0 {
-		bucket.mutex.Lock()
-		defer bucket.mutex.Unlock()
-		if exp < *bucket.nextExp || *bucket.nextExp == 0 {
-			bucket.nextExp = &exp
-			dur := expDuration(exp)
-			if dur < 0 {
-				dur = 0
-			}
-			debug("EXP: Scheduling in %s", dur)
-			if bucket.expTimer == nil {
-				bucket.expTimer = time.AfterFunc(dur, bucket.doExpiration)
-			} else {
-				bucket.expTimer.Reset(dur)
-			}
-		}
+// scheduleExpiration schedules the next expiration of documents to occur, from the minimum expiration value in the bucket. This requires locking expiration manager.
+func (bucket *Bucket) _scheduleExpiration() {
+	if nextExp, err := bucket.nextExpiration(); err == nil && nextExp > 0 {
+		bucket.expManager._scheduleExpirationAtOrBefore(nextExp)
 	}
 }
 
 func (bucket *Bucket) doExpiration() {
-	bucket.mutex.Lock()
-	bucket.nextExp = func(x uint32) *uint32 { return &x }(0)
-	bucket.mutex.Unlock()
+	bucket.expManager._clearNext()
 
 	debug("EXP: Running scheduled expiration...")
-	if n, err := bucket.ExpireDocuments(); err != nil {
-		logError("Bucket %s error expiring docs: %v", bucket, err)
+	if n, err := bucket.expireDocuments(); err != nil {
+		// If there's an error expiring docs, it means there is a programming error of a leaked expiration goroutine.
+		panic("Error expiring docs: " + err.Error())
 	} else if n > 0 {
 		info("Bucket %s expired %d docs", bucket, n)
 	}
 
-	bucket.scheduleExpiration()
+	bucket._scheduleExpiration()
 }
 
 // Completely removes all deleted documents (tombstones).

bucket_registry.go

@@ -35,18 +35,26 @@ func init() {
 	}
 }
 
-// registryBucket adds a newly opened Bucket to the registry.
-func (r *bucketRegistry) registerBucket(bucket *Bucket) {
+// registryBucket adds a newly opened Bucket to the registry. Returns true if the bucket already exists, and a copy of the bucket to use.
+func (r *bucketRegistry) registerBucket(bucket *Bucket) (bool, *Bucket) {
 	name := bucket.GetName()
 	debug("registerBucket %v %s at %s", bucket, name, bucket.url)
 	r.lock.Lock()
 	defer r.lock.Unlock()
+	return r._registerBucket(bucket)
+}
+
+// _registryBucket adds a newly opened Bucket to the registry. Returns true if the bucket already exists, and a copy of the bucket to use.
+func (r *bucketRegistry) _registerBucket(bucket *Bucket) (bool, *Bucket) {
+	name := bucket.GetName()
+	debug("_registerBucket %v %s at %s", bucket, name, bucket.url)
+
 	_, ok := r.buckets[name]
 	if !ok {
-		b := bucket.copy()
-		r.buckets[name] = b
+		r.buckets[name] = bucket
 	}
 	r.bucketCount[name] += 1
+	return ok, r.buckets[name].copy()
 }
 
 // getCachedBucket returns a bucket from the registry if it exists.
@@ -57,7 +65,9 @@ func (r *bucketRegistry) getCachedBucket(name string) *Bucket {
 	if bucket == nil {
 		return nil
 	}
-	return bucket.copy()
+	// return a copy of the bucket
+	_, bucket = r._registerBucket(bucket)
+	return bucket
 }
 
 // unregisterBucket removes a Bucket from the registry. Must be called before closing.
@@ -116,8 +126,13 @@ func getCachedBucket(name string) *Bucket {
 	return cluster.getCachedBucket(name)
 }
 
-// registryBucket adds a newly opened Bucket to the registry.
-func registerBucket(bucket *Bucket) {
+// registryBucket adds a copy of a Bucket to the registry. Returns true if the bucket already exists.
+func registerBucket(bucket *Bucket) (bool, *Bucket) {
+	return cluster.registerBucket(bucket)
+}
+
+// registryNewBucket adds a newly opened Bucket to the registry.
+func registerNewBucket(bucket *Bucket) {
 	cluster.registerBucket(bucket)
 }
 

bucket_test.go

@@ -355,7 +355,7 @@ func TestExpiration(t *testing.T) {
 	bucket := makeTestBucket(t)
 	c := bucket.DefaultDataStore()
 
-	exp, err := bucket.NextExpiration()
+	exp, err := bucket.nextExpiration()
 	require.NoError(t, err)
 	require.Equal(t, Exp(0), exp)
 
@@ -367,17 +367,17 @@ func TestExpiration(t *testing.T) {
 	requireAddRaw(t, c, "k3", 0, []byte("v3"))
 	requireAddRaw(t, c, "k4", exp4, []byte("v4"))
 
-	exp, err = bucket.NextExpiration()
+	exp, err = bucket.nextExpiration()
 	require.NoError(t, err)
 	// Usually this will return exp2, but if this is slow enough that the expiration goroutine runs to expire document k2, it can return exp4.
 	require.Contains(t, []Exp{exp2, exp4}, exp)
 
 	log.Printf("... waiting 1 sec ...")
 	time.Sleep(1 * time.Second)
 
-	exp, err = bucket.NextExpiration()
+	exp, err = bucket.nextExpiration()
 	require.NoError(t, err)
-	assert.Equal(t, exp4, exp)
+	require.Equal(t, int(exp4), int(exp))
 
 	_, _, err = c.GetRaw("k1")
 	assert.NoError(t, err)
@@ -391,7 +391,7 @@ func TestExpiration(t *testing.T) {
 	log.Printf("... waiting 2 secs ...")
 	time.Sleep(2 * time.Second)
 
-	exp, err = bucket.NextExpiration()
+	exp, err = bucket.nextExpiration()
 	require.NoError(t, err)
 	assert.Equal(t, uint32(0), exp)
 
@@ -403,6 +403,24 @@ func TestExpiration(t *testing.T) {
 	assert.Equal(t, int64(2), n)
 }
 
+func TestExpirationAfterClose(t *testing.T) {
+	t.Skip("Slow test useful for debugging issues with expiration")
+	bucket, err := OpenBucket(InMemoryURL, strings.ToLower(t.Name()), CreateNew)
+	ctx := testCtx(t)
+	defer func() {
+		assert.NoError(t, bucket.CloseAndDelete(ctx))
+	}()
+	require.NoError(t, err)
+	c := bucket.DefaultDataStore()
+
+	// set expiry long enough that Close will happen first
+	exp := Exp(time.Now().Add(1 * time.Second).Unix())
+	requireAddRaw(t, c, "docID", exp, []byte("v1"))
+	bucket.Close(ctx)
+	// sleep to ensure we won't panic
+	time.Sleep(2 * time.Second)
+}
+
 func TestUriFromPathWindows(t *testing.T) {
 	ensureNoLeaks(t)
 	if runtime.GOOS != "windows" {

collection.go

@@ -357,20 +357,23 @@ func (c *Collection) WriteCas(key string, flags int, exp Exp, cas CAS, val any,
 	return
 }
 
+// Remove creates a document tombstone. It removes the document's value and user xattrs.
 func (c *Collection) Remove(key string, cas CAS) (casOut CAS, err error) {
 	traceEnter("Remove", "%q, 0x%x", key, cas)
 	casOut, err = c.remove(key, &cas)
 	traceExit("Remove", err, "0x%x", casOut)
 	return
 }
 
+// Delete creates a document tombstone. It removes the document's value and user xattrs. Equivalent to Remove without a CAS check.
 func (c *Collection) Delete(key string) (err error) {
 	traceEnter("Delete", "%q", key)
 	_, err = c.remove(key, nil)
 	traceExit("Delete", err, "ok")
 	return err
 }
 
+// remove creates a document tombstone. It removes the document's value and user xattrs. checkClosed will allow removing the document even the bucket instance is "closed".
 func (c *Collection) remove(key string, ifCas *CAS) (casOut CAS, err error) {
 	err = c.withNewCas(func(txn *sql.Tx, newCas CAS) (e *event, err error) {
 		// Get the doc, possibly checking cas:
@@ -518,10 +521,10 @@ func (c *Collection) IsSupported(feature sgbucket.BucketStoreFeature) bool {
 
 //////// EXPIRATION
 
-// Immediately deletes all expired documents in this collection.
-func (c *Collection) ExpireDocuments() (count int64, err error) {
-	traceEnter("ExpireDocuments", "")
-	defer func() { traceExit("ExpireDocuments", err, "%d", count) }()
+// _expireDocuments immediately deletes all expired documents in this collection.
+func (c *Collection) expireDocuments() (count int64, err error) {
+	traceEnter("_expireDocuments", "")
+	defer func() { traceExit("_expireDocuments", err, "%d", count) }()
 
 	// First find all the expired docs and collect their keys:
 	exp := nowAsExpiry()