cache: support databroker option changes (#1294)

pkg/storage/inmemory/inmemory.go

@@ -52,6 +52,9 @@ type DB struct {
 	byVersion  *btree.BTree
 	deletedIDs []string
 	onchange   *signal.Signal
+
+	closeOnce sync.Once
+	closed    chan struct{}
 }
 
 // NewDB creates a new in-memory database for the given record type.
@@ -62,6 +65,7 @@ func NewDB(recordType string, btreeDegree int) *DB {
 		byID:       btree.New(btreeDegree),
 		byVersion:  btree.New(btreeDegree),
 		onchange:   s,
+		closed:     make(chan struct{}),
 	}
 }
 
@@ -84,6 +88,14 @@ func (db *DB) ClearDeleted(_ context.Context, cutoff time.Time) {
 	db.deletedIDs = remaining
 }
 
+// Close closes the database. Any watchers will be closed.
+func (db *DB) Close() error {
+	db.closeOnce.Do(func() {
+		close(db.closed)
+	})
+	return nil
+}
+
 // Delete marks a record as deleted.
 func (db *DB) Delete(_ context.Context, id string) error {
 	defer db.onchange.Broadcast()
@@ -140,7 +152,10 @@ func (db *DB) Put(_ context.Context, id string, data *anypb.Any) error {
 func (db *DB) Watch(ctx context.Context) <-chan struct{} {
 	ch := db.onchange.Bind()
 	go func() {
-		<-ctx.Done()
+		select {
+		case <-db.closed:
+		case <-ctx.Done():
+		}
 		close(ch)
 		db.onchange.Unbind(ch)
 	}()

pkg/storage/redis/redis.go

@@ -40,6 +40,9 @@ type DB struct {
 	deletedSet             string
 	tlsConfig              *tls.Config
 	notifyChMu             sync.Mutex
+
+	closeOnce sync.Once
+	closed    chan struct{}
 }
 
 // New returns new DB instance.
@@ -50,6 +53,7 @@ func New(rawURL, recordType string, deletePermanentAfter int64, opts ...Option)
 		versionSet:             recordType + "_version_set",
 		deletedSet:             recordType + "_deleted_set",
 		lastVersionKey:         recordType + "_last_version",
+		closed:                 make(chan struct{}),
 	}
 
 	for _, o := range opts {
@@ -79,6 +83,14 @@ func New(rawURL, recordType string, deletePermanentAfter int64, opts ...Option)
 	return db, nil
 }
 
+// Close closes the redis db connection.
+func (db *DB) Close() error {
+	db.closeOnce.Do(func() {
+		close(db.closed)
+	})
+	return nil
+}
+
 // Put sets new record for given id with input data.
 func (db *DB) Put(ctx context.Context, id string, data *anypb.Any) (err error) {
 	c := db.pool.Get()
@@ -259,7 +271,10 @@ func (db *DB) doNotifyLoop(ctx context.Context, ch chan struct{}) {
 	}
 	for {
 		select {
+		case <-db.closed:
+			return
 		case <-ctx.Done():
+			return
 		default:
 		}
 		switch v := psc.Receive().(type) {
@@ -271,12 +286,17 @@ func (db *DB) doNotifyLoop(ctx context.Context, ch chan struct{}) {
 			}
 
 			select {
+			case <-db.closed:
+				return
 			case <-ctx.Done():
 				log.Warn().Err(ctx.Err()).Msg("context done, stop receive from redis channel")
 				return
 			default:
 				db.notifyChMu.Lock()
 				select {
+				case <-db.closed:
+					db.notifyChMu.Unlock()
+					return
 				case <-ctx.Done():
 					db.notifyChMu.Unlock()
 					log.Warn().Err(ctx.Err()).Msg("context done while holding notify lock, stop receive from redis channel")
@@ -313,6 +333,7 @@ func (db *DB) watch(ctx context.Context, ch chan struct{}) {
 		db.doNotifyLoop(ctx, ch)
 	}()
 	select {
+	case <-db.closed:
 	case <-ctx.Done():
 	case <-done:
 	}

pkg/storage/storage.go

@@ -12,6 +12,9 @@ import (
 
 // Backend is the interface required for a storage backend.
 type Backend interface {
+	// Close closes the backend.
+	Close() error
+
 	// Put is used to insert or update a record.
 	Put(ctx context.Context, id string, data *anypb.Any) error
 

pkg/storage/storage_test.go

@@ -4,8 +4,9 @@ import (
 	"context"
 	"time"
 
-	"github.com/pomerium/pomerium/pkg/grpc/databroker"
 	"google.golang.org/protobuf/types/known/anypb"
+
+	"github.com/pomerium/pomerium/pkg/grpc/databroker"
 )
 
 type mockBackend struct {
@@ -18,6 +19,10 @@ type mockBackend struct {
 	watch        func(ctx context.Context) <-chan struct{}
 }
 
+func (m *mockBackend) Close() error {
+	return nil
+}
+
 func (m *mockBackend) Put(ctx context.Context, id string, data *anypb.Any) error {
 	return m.put(ctx, id, data)
 }