github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/internal/datastore/memdb/watch.go (about)

     1  package memdb
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"fmt"
     7  	"time"
     8  
     9  	"github.com/hashicorp/go-memdb"
    10  
    11  	"github.com/authzed/spicedb/internal/datastore/revisions"
    12  	"github.com/authzed/spicedb/pkg/datastore"
    13  )
    14  
    15  const errWatchError = "watch error: %w"
    16  
    17  func (mdb *memdbDatastore) Watch(ctx context.Context, ar datastore.Revision, options datastore.WatchOptions) (<-chan *datastore.RevisionChanges, <-chan error) {
    18  	watchBufferLength := options.WatchBufferLength
    19  	if watchBufferLength == 0 {
    20  		watchBufferLength = mdb.watchBufferLength
    21  	}
    22  
    23  	updates := make(chan *datastore.RevisionChanges, watchBufferLength)
    24  	errs := make(chan error, 1)
    25  
    26  	watchBufferWriteTimeout := options.WatchBufferWriteTimeout
    27  	if watchBufferWriteTimeout == 0 {
    28  		watchBufferWriteTimeout = mdb.watchBufferWriteTimeout
    29  	}
    30  
    31  	sendChange := func(change *datastore.RevisionChanges) bool {
    32  		select {
    33  		case updates <- change:
    34  			return true
    35  
    36  		default:
    37  			// If we cannot immediately write, setup the timer and try again.
    38  		}
    39  
    40  		timer := time.NewTimer(watchBufferWriteTimeout)
    41  		defer timer.Stop()
    42  
    43  		select {
    44  		case updates <- change:
    45  			return true
    46  
    47  		case <-timer.C:
    48  			errs <- datastore.NewWatchDisconnectedErr()
    49  			return false
    50  		}
    51  	}
    52  
    53  	go func() {
    54  		defer close(updates)
    55  		defer close(errs)
    56  
    57  		currentTxn := ar.(revisions.TimestampRevision).TimestampNanoSec()
    58  
    59  		for {
    60  			var stagedUpdates []*datastore.RevisionChanges
    61  			var watchChan <-chan struct{}
    62  			var err error
    63  			stagedUpdates, currentTxn, watchChan, err = mdb.loadChanges(ctx, currentTxn, options)
    64  			if err != nil {
    65  				errs <- err
    66  				return
    67  			}
    68  
    69  			// Write the staged updates to the channel
    70  			for _, changeToWrite := range stagedUpdates {
    71  				if !sendChange(changeToWrite) {
    72  					return
    73  				}
    74  			}
    75  
    76  			// Wait for new changes
    77  			ws := memdb.NewWatchSet()
    78  			ws.Add(watchChan)
    79  
    80  			err = ws.WatchCtx(ctx)
    81  			if err != nil {
    82  				switch {
    83  				case errors.Is(err, context.Canceled):
    84  					errs <- datastore.NewWatchCanceledErr()
    85  				default:
    86  					errs <- fmt.Errorf(errWatchError, err)
    87  				}
    88  				return
    89  			}
    90  		}
    91  	}()
    92  
    93  	return updates, errs
    94  }
    95  
    96  func (mdb *memdbDatastore) loadChanges(_ context.Context, currentTxn int64, options datastore.WatchOptions) ([]*datastore.RevisionChanges, int64, <-chan struct{}, error) {
    97  	mdb.RLock()
    98  	defer mdb.RUnlock()
    99  
   100  	loadNewTxn := mdb.db.Txn(false)
   101  	defer loadNewTxn.Abort()
   102  
   103  	it, err := loadNewTxn.LowerBound(tableChangelog, indexRevision, currentTxn+1)
   104  	if err != nil {
   105  		return nil, 0, nil, fmt.Errorf(errWatchError, err)
   106  	}
   107  
   108  	var changes []*datastore.RevisionChanges
   109  	lastRevision := currentTxn
   110  	for changeRaw := it.Next(); changeRaw != nil; changeRaw = it.Next() {
   111  		change := changeRaw.(*changelog)
   112  
   113  		if options.Content&datastore.WatchRelationships == datastore.WatchRelationships && len(change.changes.RelationshipChanges) > 0 {
   114  			changes = append(changes, &change.changes)
   115  		}
   116  
   117  		if options.Content&datastore.WatchCheckpoints == datastore.WatchCheckpoints && change.revisionNanos > lastRevision {
   118  			changes = append(changes, &datastore.RevisionChanges{
   119  				Revision:     revisions.NewForTimestamp(change.revisionNanos),
   120  				IsCheckpoint: true,
   121  			})
   122  		}
   123  
   124  		if options.Content&datastore.WatchSchema == datastore.WatchSchema &&
   125  			len(change.changes.ChangedDefinitions) > 0 || len(change.changes.DeletedCaveats) > 0 || len(change.changes.DeletedNamespaces) > 0 {
   126  			changes = append(changes, &change.changes)
   127  		}
   128  
   129  		lastRevision = change.revisionNanos
   130  	}
   131  
   132  	watchChan, _, err := loadNewTxn.LastWatch(tableChangelog, indexRevision)
   133  	if err != nil {
   134  		return nil, 0, nil, fmt.Errorf(errWatchError, err)
   135  	}
   136  
   137  	return changes, lastRevision, watchChan, nil
   138  }