Add new API Subscribe to listen on update for the specified prefix (#771)

Signed-off-by: பாலாஜி <rbalajis25@gmail.com>

Author: poonai

badger/cmd/bank.go

@@ -68,6 +68,7 @@ var duration string
 var stopAll int32
 var mmap bool
 var checkStream bool
+var checkSubscriber bool
 
 const keyPrefix = "account:"
 
@@ -87,6 +88,9 @@ func init() {
 	bankTest.Flags().BoolVarP(&checkStream, "check_stream", "s", false,
 		"If true, the test will send transactions to another badger instance via the stream "+
 			"interface in order to verify that all data is streamed correctly.")
+	bankTest.Flags().BoolVarP(&checkSubscriber, "check_subscriber", "w", false,
+		"If true, the test will send transactions to another badger instance via the subscriber "+
+			"interface in order to verify that all the data is published correctly.")
 
 	bankDisect.Flags().IntVarP(&numPrevious, "previous", "p", 12,
 		"Starting from the violation txn, how many previous versions to retrieve.")
@@ -347,6 +351,24 @@ func runTest(cmd *cobra.Command, args []string) error {
 	defer db.Close()
 
 	var tmpDb *badger.DB
+	var subscribeDB *badger.DB
+	if checkSubscriber {
+		dir, err := ioutil.TempDir("", "bank_subscribe")
+		y.Check(err)
+
+		subscribeOpts := badger.DefaultOptions
+		subscribeOpts.Dir = dir
+		subscribeOpts.ValueDir = dir
+		subscribeOpts.SyncWrites = false
+		log.Printf("Opening subscribe DB with options: %+v\n", subscribeOpts)
+
+		subscribeDB, err = badger.Open(subscribeOpts)
+		if err != nil {
+			return err
+		}
+		defer subscribeDB.Close()
+	}
+
 	if checkStream {
 		dir, err := ioutil.TempDir("", "bank_stream")
 		y.Check(err)
@@ -511,8 +533,39 @@ func runTest(cmd *cobra.Command, args []string) error {
 			}))
 		}
 	}()
-	wg.Wait()
 
+	ctx, cancel := context.WithCancel(context.Background())
+	var subWg sync.WaitGroup
+	if checkSubscriber {
+		subWg.Add(1)
+		go func() {
+			defer subWg.Done()
+			accountIDS := [][]byte{}
+			for i := 0; i < numAccounts; i++ {
+				accountIDS = append(accountIDS, key(i))
+			}
+			updater := func(kvs *pb.KVList) {
+				batch := subscribeDB.NewWriteBatch()
+				for _, kv := range kvs.GetKv() {
+					y.Check(batch.Set(kv.Key, kv.Value, 0))
+				}
+				y.Check(batch.Flush())
+			}
+			db.Subscribe(ctx, updater, accountIDS[0], accountIDS[1:]...)
+		}()
+	}
+	wg.Wait()
+	cancel()
+	subWg.Wait()
+	y.Check(subscribeDB.View(func(txn *badger.Txn) error {
+		_, err := seekTotal(txn)
+		if err != nil {
+			log.Printf("Error while calculating subscriber DB total: %v", err)
+		} else {
+			atomic.AddUint64(&reads, 1)
+		}
+		return nil
+	}))
 	if atomic.LoadInt32(&stopAll) == 0 {
 		log.Println("Test OK")
 		return nil

db.go

@@ -18,6 +18,7 @@ package badger
 
 import (
 	"bytes"
+	"context"
 	"encoding/binary"
 	"encoding/hex"
 	"expvar"
@@ -32,6 +33,7 @@ import (
 	"time"
 
 	"github.com/dgraph-io/badger/options"
+	"github.com/dgraph-io/badger/pb"
 	"github.com/dgraph-io/badger/skl"
 	"github.com/dgraph-io/badger/table"
 	"github.com/dgraph-io/badger/y"
@@ -54,8 +56,11 @@ type closers struct {
 	memtable   *y.Closer
 	writes     *y.Closer
 	valueGC    *y.Closer
+	pub        *y.Closer
 }
 
+type callback func(kv *pb.KVList)
+
 // DB provides the various functions required to interact with Badger.
 // DB is thread-safe.
 type DB struct {
@@ -85,6 +90,8 @@ type DB struct {
 	blockWrites int32
 
 	orc *oracle
+
+	pub *publisher
 }
 
 const (
@@ -267,6 +274,7 @@ func Open(opt Options) (db *DB, err error) {
 		dirLockGuard:  dirLockGuard,
 		valueDirGuard: valueDirLockGuard,
 		orc:           newOracle(opt),
+		pub:           newPublisher(),
 	}
 
 	// Calculate initial size.
@@ -323,6 +331,9 @@ func Open(opt Options) (db *DB, err error) {
 	db.closers.valueGC = y.NewCloser(1)
 	go db.vlog.waitOnGC(db.closers.valueGC)
 
+	db.closers.pub = y.NewCloser(1)
+	go db.pub.listenForUpdates(db.closers.pub)
+
 	valueDirLockGuard = nil
 	dirLockGuard = nil
 	manifestFile = nil
@@ -342,6 +353,8 @@ func (db *DB) Close() (err error) {
 	// Stop writes next.
 	db.closers.writes.SignalAndWait()
 
+	db.closers.pub.SignalAndWait()
+
 	// Now close the value log.
 	if vlogErr := db.vlog.Close(); vlogErr != nil {
 		err = errors.Wrap(vlogErr, "DB.Close")
@@ -610,6 +623,8 @@ func (db *DB) writeRequests(reqs []*request) error {
 		return err
 	}
 
+	db.elog.Printf("Sending updates to subscribers")
+	db.pub.sendUpdates(reqs)
 	db.elog.Printf("Writing to memtable")
 	var count int
 	for _, b := range reqs {
@@ -662,6 +677,8 @@ func (db *DB) sendToWriteCh(entries []*Entry) (*request, error) {
 	req.Entries = entries
 	req.Wg = sync.WaitGroup{}
 	req.Wg.Add(1)
+	req.IncrRef()     // for db write
+	req.IncrRef()     // for publisher updates
 	db.writeCh <- req // Handled in doWrites.
 	y.NumPuts.Add(int64(len(entries)))
 
@@ -1409,3 +1426,46 @@ func (db *DB) DropPrefix(prefix []byte) error {
 	db.opt.Infof("DropPrefix done")
 	return nil
 }
+
+// Subscribe can be used watch key changes for the given key prefix.
+func (db *DB) Subscribe(ctx context.Context, cb callback, prefix []byte, prefixes ...[]byte) error {
+	if cb == nil {
+		return ErrNilCallback
+	}
+	prefixes = append(prefixes, prefix)
+	c := y.NewCloser(1)
+	recvCh, id := db.pub.newSubscriber(c, prefixes...)
+	slurp := func(batch *pb.KVList) {
+		defer func() {
+			if len(batch.GetKv()) > 0 {
+				cb(batch)
+			}
+		}()
+		for {
+			select {
+			case kvs := <-recvCh:
+				batch.Kv = append(batch.Kv, kvs.Kv...)
+			default:
+				return
+			}
+		}
+	}
+	for {
+		select {
+		case <-c.HasBeenClosed():
+			slurp(new(pb.KVList))
+			// Drain if any pending updates.
+			c.Done()
+			// No need to delete here. Closer will be called only while
+			// closing DB. Subscriber will be deleted by cleanSubscribers.
+			return nil
+		case <-ctx.Done():
+			c.Done()
+			db.pub.deleteSubscriber(id)
+			// Delete the subscriber to avoid further updates.
+			return ctx.Err()
+		case batch := <-recvCh:
+			slurp(batch)
+		}
+	}
+}

db_test.go

@@ -18,6 +18,7 @@ package badger
 
 import (
 	"bytes"
+	"context"
 	"encoding/binary"
 	"flag"
 	"fmt"
@@ -35,6 +36,7 @@ import (
 	"time"
 
 	"github.com/dgraph-io/badger/options"
+	"github.com/dgraph-io/badger/pb"
 	"github.com/dgraph-io/badger/skl"
 
 	"github.com/dgraph-io/badger/y"
@@ -1514,10 +1516,31 @@ func TestGoroutineLeak(t *testing.T) {
 	t.Logf("Num go: %d", before)
 	for i := 0; i < 12; i++ {
 		runBadgerTest(t, nil, func(t *testing.T, db *DB) {
+			updated := false
+			ctx, cancel := context.WithCancel(context.Background())
+			var wg sync.WaitGroup
+			wg.Add(1)
+			var subWg sync.WaitGroup
+			subWg.Add(1)
+			go func() {
+				subWg.Done()
+				err := db.Subscribe(ctx, func(kvs *pb.KVList) {
+					require.Equal(t, []byte("value"), kvs.Kv[0].GetValue())
+					updated = true
+					wg.Done()
+				}, []byte("key"))
+				if err != nil {
+					require.Equal(t, err.Error(), context.Canceled.Error())
+				}
+			}()
+			subWg.Wait()
 			err := db.Update(func(txn *Txn) error {
 				return txn.Set([]byte("key"), []byte("value"))
 			})
 			require.NoError(t, err)
+			wg.Wait()
+			cancel()
+			require.Equal(t, true, updated)
 		})
 	}
 	require.Equal(t, before, runtime.NumGoroutine())

errors.go

@@ -102,4 +102,7 @@ var (
 	// ErrBlockedWrites is returned if the user called DropAll. During the process of dropping all
 	// data from Badger, we stop accepting new writes, by returning this error.
 	ErrBlockedWrites = errors.New("Writes are blocked, possibly due to DropAll or Close")
+
+	// ErrNilCallback is returned when subscriber's callback is nil.
+	ErrNilCallback = errors.New("Callback cannot be nil")
 )