Refactor Store interface to eliminate IDs for Trigger registration

The store used to require the user to supply a unique ID in order to
register a trigger. Instead, when registering a trigger, the Store now
returns a handle on which the trigger can be unregistered. This has a
number of benefits, including that the user no longer has to reason
about what IDs they have already used, and only the owner of the handle
can end the trigger.

server/server.go

@@ -131,7 +131,7 @@ func (s *Server) Init(ctx context.Context) (*Server, error) {
 	config := storage.TriggerConfig{
 		OnCommit: s.reload,
 	}
-	if err := s.store.Register(ctx, txn, "opa/server", config); err != nil {
+	if _, err := s.store.Register(ctx, txn, config); err != nil {
 		s.store.Abort(ctx, txn)
 		return nil, err
 	}

server/server_test.go

@@ -1050,8 +1050,8 @@ func (queryBindingErrStore) Abort(ctx context.Context, txn storage.Transaction)
 
 }
 
-func (queryBindingErrStore) Register(context.Context, storage.Transaction, string, storage.TriggerConfig) error {
-	return nil
+func (queryBindingErrStore) Register(context.Context, storage.Transaction, storage.TriggerConfig) (storage.TriggerHandle, error) {
+	return nil, nil
 }
 
 func (queryBindingErrStore) Unregister(context.Context, storage.Transaction, string) {

storage/inmem/inmem.go

@@ -31,7 +31,7 @@ import (
 func New() storage.Store {
 	return &store{
 		data:     map[string]interface{}{},
-		triggers: map[string]storage.TriggerConfig{},
+		triggers: map[*handle]storage.TriggerConfig{},
 		policies: map[string][]byte{},
 		indices:  newIndices(),
 	}
@@ -66,13 +66,17 @@ func NewFromReader(r io.Reader) storage.Store {
 }
 
 type store struct {
-	rmu      sync.RWMutex                     // reader-writer lock
-	wmu      sync.Mutex                       // writer lock
-	xid      uint64                           // last generated transaction id
-	data     map[string]interface{}           // raw data
-	policies map[string][]byte                // raw policies
-	triggers map[string]storage.TriggerConfig // registered triggers
-	indices  *indices                         // data ref indices
+	rmu      sync.RWMutex                      // reader-writer lock
+	wmu      sync.Mutex                        // writer lock
+	xid      uint64                            // last generated transaction id
+	data     map[string]interface{}            // raw data
+	policies map[string][]byte                 // raw policies
+	triggers map[*handle]storage.TriggerConfig // registered triggers
+	indices  *indices                          // data ref indices
+}
+
+type handle struct {
+	db *store
 }
 
 func (db *store) NewTransaction(ctx context.Context, params ...storage.TransactionParams) (storage.Transaction, error) {
@@ -136,24 +140,17 @@ func (db *store) DeletePolicy(_ context.Context, txn storage.Transaction, id str
 	return underlying.DeletePolicy(id)
 }
 
-func (db *store) Register(ctx context.Context, txn storage.Transaction, id string, config storage.TriggerConfig) error {
+func (db *store) Register(ctx context.Context, txn storage.Transaction, config storage.TriggerConfig) (storage.TriggerHandle, error) {
 	underlying := txn.(*transaction)
 	if !underlying.write {
-		return &storage.Error{
+		return nil, &storage.Error{
 			Code:    storage.InvalidTransactionErr,
 			Message: "triggers must be registered with a write transaction",
 		}
 	}
-	db.triggers[id] = config
-	return nil
-}
-
-func (db *store) Unregister(ctx context.Context, txn storage.Transaction, id string) {
-	underlying := txn.(*transaction)
-	if !underlying.write {
-		return
-	}
-	delete(db.triggers, id)
+	h := &handle{db}
+	db.triggers[h] = config
+	return h, nil
 }
 
 func (db *store) Read(ctx context.Context, txn storage.Transaction, path storage.Path) (interface{}, error) {
@@ -177,6 +174,14 @@ func (db *store) Build(ctx context.Context, txn storage.Transaction, ref ast.Ref
 	return db.indices.Build(ctx, db, txn, ref)
 }
 
+func (h *handle) Unregister(ctx context.Context, txn storage.Transaction) {
+	underlying := txn.(*transaction)
+	if !underlying.write {
+		return
+	}
+	delete(h.db.triggers, h)
+}
+
 func (db *store) runOnCommitTriggers(ctx context.Context, txn storage.Transaction, event storage.TriggerEvent) {
 	for _, t := range db.triggers {
 		t.OnCommit(ctx, txn, event)

storage/inmem/inmem_test.go

@@ -438,7 +438,7 @@ func TestInMemoryTriggers(t *testing.T) {
 	writeTxn := storage.NewTransactionOrDie(ctx, store, storage.WriteParams)
 	readTxn := storage.NewTransactionOrDie(ctx, store)
 
-	err := store.Register(ctx, readTxn, "err", storage.TriggerConfig{
+	_, err := store.Register(ctx, readTxn, storage.TriggerConfig{
 		OnCommit: func(context.Context, storage.Transaction, storage.TriggerEvent) {},
 	})
 
@@ -452,7 +452,7 @@ func TestInMemoryTriggers(t *testing.T) {
 	modifiedPath := storage.MustParsePath("/a")
 	expectedValue := "hello"
 
-	err = store.Register(ctx, writeTxn, "test", storage.TriggerConfig{
+	_, err = store.Register(ctx, writeTxn, storage.TriggerConfig{
 		OnCommit: func(ctx context.Context, txn storage.Transaction, evt storage.TriggerEvent) {
 			result, err := store.Read(ctx, txn, modifiedPath)
 			if err != nil || !reflect.DeepEqual(result, expectedValue) {
@@ -461,6 +461,9 @@ func TestInMemoryTriggers(t *testing.T) {
 			event = evt
 		},
 	})
+	if err != nil {
+		t.Fatalf("Failed to register callback: %v", err)
+	}
 
 	if err := store.Write(ctx, writeTxn, storage.ReplaceOp, modifiedPath, expectedValue); err != nil {
 		t.Fatalf("Unexpected write error: %v", err)
@@ -491,6 +494,55 @@ func TestInMemoryTriggers(t *testing.T) {
 	}
 }
 
+func TestInMemoryTriggersUnregister(t *testing.T) {
+	ctx := context.Background()
+	store := NewFromObject(loadSmallTestData())
+	writeTxn := storage.NewTransactionOrDie(ctx, store, storage.WriteParams)
+	modifiedPath := storage.MustParsePath("/a")
+	expectedValue := "hello"
+
+	var called bool
+	_, err := store.Register(ctx, writeTxn, storage.TriggerConfig{
+		OnCommit: func(ctx context.Context, txn storage.Transaction, evt storage.TriggerEvent) {
+			if !evt.IsZero() {
+				called = true
+			}
+		},
+	})
+	if err != nil {
+		t.Fatalf("Failed to register callback: %v", err)
+	}
+
+	handle, err := store.Register(ctx, writeTxn, storage.TriggerConfig{
+		OnCommit: func(ctx context.Context, txn storage.Transaction, evt storage.TriggerEvent) {
+			if !evt.IsZero() {
+				t.Fatalf("Callback should have been unregistered")
+			}
+		},
+	})
+	if err != nil {
+		t.Fatalf("Failed to register callback: %v", err)
+	}
+
+	if err := store.Commit(ctx, writeTxn); err != nil {
+		t.Fatalf("Unexpected commit error: %v", err)
+	}
+
+	writeTxn = storage.NewTransactionOrDie(ctx, store, storage.WriteParams)
+	if err := store.Write(ctx, writeTxn, storage.AddOp, modifiedPath, expectedValue); err != nil {
+		t.Fatalf("Failed to write to store: %v", err)
+	}
+	handle.Unregister(ctx, writeTxn)
+
+	if err := store.Commit(ctx, writeTxn); err != nil {
+		t.Fatalf("Unexpected commit error: %v", err)
+	}
+
+	if !called {
+		t.Fatal("Registered callback was not called")
+	}
+}
+
 func loadExpectedResult(input string) interface{} {
 	if len(input) == 0 {
 		return nil

storage/interface.go

@@ -149,8 +149,7 @@ type TriggerConfig struct {
 // Trigger defines the interface that stores implement to register for change
 // notifications when the store is changed.
 type Trigger interface {
-	Register(ctx context.Context, txn Transaction, id string, config TriggerConfig) error
-	Unregister(ctx context.Context, txn Transaction, id string)
+	Register(ctx context.Context, txn Transaction, config TriggerConfig) (TriggerHandle, error)
 }
 
 // TriggersNotSupported provides default implementations of the Trigger
@@ -158,12 +157,14 @@ type Trigger interface {
 type TriggersNotSupported struct{}
 
 // Register always returns an error indicating triggers are not supported.
-func (TriggersNotSupported) Register(context.Context, Transaction, string, TriggerConfig) error {
-	return triggersNotSupportedError()
+func (TriggersNotSupported) Register(context.Context, Transaction, TriggerConfig) (TriggerHandle, error) {
+	return nil, triggersNotSupportedError()
 }
 
-// Unregister is a no-op.
-func (TriggersNotSupported) Unregister(context.Context, Transaction, string) {
+// TriggerHandle defines the interface that can be used to unregister triggers that have
+// been registered on a Store.
+type TriggerHandle interface {
+	Unregister(ctx context.Context, txn Transaction)
 }
 
 // IndexIterator defines the interface for iterating over index results.