Stream the full set of predicates and types during a snapshot. (hypermodeinc#5444)

Type and predicate keys always have a timestamp of 1 so they are not being picked up by every snapshot.

This PR adds logic to send the entire schema in the snapshot and to delete predicates and types
not present in the snapshot (since they were deleted in between snapshots).

Author: martinmr

go.mod

@@ -16,7 +16,7 @@ require (
 	github.com/blevesearch/segment v0.0.0-20160915185041-762005e7a34f // indirect
 	github.com/blevesearch/snowballstem v0.0.0-20180110192139-26b06a2c243d // indirect
 	github.com/codahale/hdrhistogram v0.0.0-20161010025455-3a0bb77429bd
-	github.com/dgraph-io/badger/v2 v2.0.1-rc1.0.20200506081535-536fed1846d0
+	github.com/dgraph-io/badger/v2 v2.0.1-rc1.0.20200515210839-ef28ef36b592
 	github.com/dgraph-io/dgo/v200 v200.0.0-20200401175452-e463f9234453
 	github.com/dgraph-io/ristretto v0.0.2
 	github.com/dgrijalva/jwt-go v3.2.0+incompatible

go.sum

@@ -76,8 +76,12 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/dgraph-io/badger v1.6.0 h1:DshxFxZWXUcO0xX476VJC07Xsr6ZCBVRHKZ93Oh7Evo=
 github.com/dgraph-io/badger v1.6.0/go.mod h1:zwt7syl517jmP8s94KqSxTlM6IMsdhYy6psNgSztDR4=
-github.com/dgraph-io/badger/v2 v2.0.1-rc1.0.20200506081535-536fed1846d0 h1:4VBIyLibX6qFfz6wSbEvp4RBfoKETvHfIln18ROLiHI=
-github.com/dgraph-io/badger/v2 v2.0.1-rc1.0.20200506081535-536fed1846d0/go.mod h1:3KY8+bsP8wI0OEnQJAKpd4wIJW/Mm32yw2j/9FUVnIM=
+github.com/dgraph-io/badger/v2 v2.0.1-rc1.0.20200515084818-8ee988574454 h1:xxDt2YtzZ2tVCWVk4tcYyKtZVtWRcScAcsbNhHcE/Qg=
+github.com/dgraph-io/badger/v2 v2.0.1-rc1.0.20200515084818-8ee988574454/go.mod h1:3KY8+bsP8wI0OEnQJAKpd4wIJW/Mm32yw2j/9FUVnIM=
+github.com/dgraph-io/badger/v2 v2.0.1-rc1.0.20200515182826-fafb627b4ce4 h1:0hhuliB6NlxxmqQdrl0MkZi+qjibu69uPKtckRn8E0I=
+github.com/dgraph-io/badger/v2 v2.0.1-rc1.0.20200515182826-fafb627b4ce4/go.mod h1:3KY8+bsP8wI0OEnQJAKpd4wIJW/Mm32yw2j/9FUVnIM=
+github.com/dgraph-io/badger/v2 v2.0.1-rc1.0.20200515210839-ef28ef36b592 h1:j6j3yunDbktI4H3tbj3grt2enO4EPbhstU6Tb8HwqdQ=
+github.com/dgraph-io/badger/v2 v2.0.1-rc1.0.20200515210839-ef28ef36b592/go.mod h1:3KY8+bsP8wI0OEnQJAKpd4wIJW/Mm32yw2j/9FUVnIM=
 github.com/dgraph-io/dgo/v200 v200.0.0-20200401175452-e463f9234453 h1:DTgOrw91nMIukDm/WEvdobPLl0LgeDd/JE66+24jBks=
 github.com/dgraph-io/dgo/v200 v200.0.0-20200401175452-e463f9234453/go.mod h1:Co+FwJrnndSrPORO8Gdn20dR7FPTfmXr0W/su0Ve/Ig=
 github.com/dgraph-io/ristretto v0.0.2-0.20200115201040-8f368f2f2ab3 h1:MQLRM35Pp0yAyBYksjbj1nZI/w6eyRY/mWoM1sFf4kU=

protos/pb.proto

@@ -295,9 +295,14 @@ message Proposal {
 }
 
 message KVS {
- repeated badgerpb2.KV kv = 1;
- // done used to indicate if the stream of KVS is over.
- bool done      = 2;
+	repeated badgerpb2.KV kv = 1;
+
+	// done used to indicate if the stream of KVS is over.
+	bool done	= 2;
+	// predicates is the list of predicates known by the leader at the time of the snapshot.
+	repeated string predicates = 3;
+	// types is the list of types known by the leader at the time of the snapshot.
+	repeated string types = 4;
 }
 
 // Posting messages.

protos/pb/pb.pb.go

@@ -117,6 +117,10 @@ func (s *state) Delete(attr string) error {
 
 // DeleteType updates the schema in memory and disk
 func (s *state) DeleteType(typeName string) error {
+	if s == nil {
+		return nil
+	}
+
 	s.Lock()
 	defer s.Unlock()
 
@@ -265,6 +269,10 @@ func (s *state) IsIndexed(ctx context.Context, pred string) bool {
 
 // Predicates returns the list of predicates for given group
 func (s *state) Predicates() []string {
+	if s == nil {
+		return nil
+	}
+
 	s.RLock()
 	defer s.RUnlock()
 	var out []string
@@ -276,6 +284,10 @@ func (s *state) Predicates() []string {
 
 // Types returns the list of types.
 func (s *state) Types() []string {
+	if s == nil {
+		return nil
+	}
+
 	s.RLock()
 	defer s.RUnlock()
 	var out []string

schema/schema.go

@@ -173,6 +173,23 @@ func RetryQuery(dg *dgo.Dgraph, q string) (*api.Response, error) {
 	}
 }
 
+// RetryBadQuery will retry a query until it failse with a non-retryable error.
+func RetryBadQuery(dg *dgo.Dgraph, q string) (*api.Response, error) {
+	for {
+		txn := dg.NewTxn()
+		ctx := context.Background()
+		resp, err := txn.Query(ctx, q)
+		if err == nil || strings.Contains(err.Error(), "Please retry") {
+			time.Sleep(10 * time.Millisecond)
+			txn.Discard(ctx)
+			continue
+		}
+
+		txn.Discard(ctx)
+		return resp, err
+	}
+}
+
 // RetryMutation will retry a mutation until it succeeds or a non-retryable error is received.
 // The mutation should have CommitNow set to true.
 func RetryMutation(dg *dgo.Dgraph, mu *api.Mutation) error {

testutil/client.go

@@ -17,16 +17,21 @@
 package worker
 
 import (
+	"context"
 	"sync/atomic"
+	"time"
 
 	bpb "github.com/dgraph-io/badger/v2/pb"
 	"github.com/golang/glog"
+	"github.com/pkg/errors"
 	"go.etcd.io/etcd/raft"
 
 	"github.com/dgraph-io/badger/v2"
 	"github.com/dgraph-io/dgraph/conn"
 	"github.com/dgraph-io/dgraph/posting"
 	"github.com/dgraph-io/dgraph/protos/pb"
+	"github.com/dgraph-io/dgraph/schema"
+	"github.com/dgraph-io/dgraph/x"
 )
 
 const (
@@ -70,12 +75,14 @@ func (n *node) populateSnapshot(snap pb.Snapshot, pl *conn.Pool) (int, error) {
 
 	// We can use count to check the number of posting lists returned in tests.
 	count := 0
+	var done *pb.KVS
 	for {
 		kvs, err := stream.Recv()
 		if err != nil {
 			return count, err
 		}
 		if kvs.Done {
+			done = kvs
 			glog.V(1).Infoln("All key-values have been received.")
 			break
 		}
@@ -95,6 +102,10 @@ func (n *node) populateSnapshot(snap pb.Snapshot, pl *conn.Pool) (int, error) {
 		return 0, err
 	}
 
+	if err := deleteStalePreds(ctx, done); err != nil {
+		return count, err
+	}
+
 	glog.Infof("Snapshot writes DONE. Sending ACK")
 	// Send an acknowledgement back to the leader.
 	if err := stream.Send(&pb.Snapshot{Done: true}); err != nil {
@@ -104,6 +115,60 @@ func (n *node) populateSnapshot(snap pb.Snapshot, pl *conn.Pool) (int, error) {
 	return count, nil
 }
 
+func deleteStalePreds(ctx context.Context, kvs *pb.KVS) error {
+	if kvs == nil {
+		return nil
+	}
+
+	// Look for predicates present in the receiver but not in the list sent by the leader.
+	// These predicates were deleted in between snapshots and need to be deleted from the
+	// receiver to keep the schema in sync.
+	currPredicates := schema.State().Predicates()
+	snapshotPreds := make(map[string]struct{})
+	for _, pred := range kvs.Predicates {
+		snapshotPreds[pred] = struct{}{}
+	}
+	for _, pred := range currPredicates {
+		if _, ok := snapshotPreds[pred]; !ok {
+		LOOP:
+			for {
+				err := posting.DeletePredicate(ctx, pred)
+				switch err {
+				case badger.ErrBlockedWrites:
+					time.Sleep(1 * time.Second)
+				case nil:
+					break LOOP
+				default:
+					glog.Warningf(
+						"Cannot delete removed predicate %s after streaming snapshot: %v",
+						pred, err)
+					return errors.Wrapf(err,
+						"cannot delete removed predicate %s after streaming snapshot", pred)
+				}
+			}
+		}
+	}
+
+	// Look for types present in the receiver but not in the list sent by the leader.
+	// These types were deleted in between snapshots and need to be deleted from the
+	// receiver to keep the schema in sync.
+	currTypes := schema.State().Types()
+	snapshotTypes := make(map[string]struct{})
+	for _, typ := range kvs.Types {
+		snapshotTypes[typ] = struct{}{}
+	}
+	for _, typ := range currTypes {
+		if _, ok := snapshotTypes[typ]; !ok {
+			if err := schema.State().DeleteType(typ); err != nil {
+				return errors.Wrapf(err, "cannot delete removed type %s after streaming snapshot",
+					typ)
+			}
+		}
+	}
+
+	return nil
+}
+
 func doStreamSnapshot(snap *pb.Snapshot, out pb.Worker_StreamSnapshotServer) error {
 	// We choose not to try and match the requested snapshot from the latest snapshot at the leader.
 	// This is the job of the Raft library. At the leader end, we service whatever is asked of us.
@@ -138,15 +203,35 @@ func doStreamSnapshot(snap *pb.Snapshot, out pb.Worker_StreamSnapshotServer) err
 		return out.Send(kvs)
 	}
 	stream.ChooseKey = func(item *badger.Item) bool {
-		return item.Version() >= snap.SinceTs
+		if item.Version() >= snap.SinceTs {
+			return true
+		}
+
+		if item.Version() != 1 {
+			return false
+		}
+
+		// Type and Schema keys always have a timestamp of 1. They all need to be sent
+		// with the snapshot.
+		pk, err := x.Parse(item.Key())
+		if err != nil {
+			return false
+		}
+		return pk.IsSchema() || pk.IsType()
 	}
 
 	if err := stream.Orchestrate(out.Context()); err != nil {
 		return err
 	}
 
-	// Indicate that sending is done.
-	if err := out.Send(&pb.KVS{Done: true}); err != nil {
+	// Indicate that sending is done. Send a list of all the predicate and types at the
+	// time of the snapshot so that the receiver can delete predicates
+	done := &pb.KVS{
+		Done:       true,
+		Predicates: schema.State().Predicates(),
+		Types:      schema.State().Types(),
+	}
+	if err := out.Send(done); err != nil {
 		return err
 	}
 

worker/snapshot.go

@@ -44,12 +44,26 @@ func TestSnapshot(t *testing.T) {
 		DropOp: api.Operation_ALL,
 	}))
 	require.NoError(t, dg1.Alter(context.Background(), &api.Operation{
-		Schema: "value: int .",
+		Schema: `
+			value: int .
+			name: string .
+			address: string @index(term) .`,
 	}))
 
 	err = testutil.DockerStop("alpha2")
 	require.NoError(t, err)
 
+	// Update the name predicate to include an index.
+	require.NoError(t, dg1.Alter(context.Background(), &api.Operation{
+		Schema: `name: string @index(term) .`,
+	}))
+
+	// Delete the address predicate.
+	require.NoError(t, dg1.Alter(context.Background(), &api.Operation{
+		DropOp:    api.Operation_ATTR,
+		DropValue: "address",
+	}))
+
 	for i := 1; i <= 200; i++ {
 		err := testutil.RetryMutation(dg1, &api.Mutation{
 			SetNquads: []byte(fmt.Sprintf(`_:node <value> "%d" .`, i)),
@@ -112,6 +126,30 @@ func verifySnapshot(t *testing.T, dg *dgo.Dgraph, num int) {
 		sum += item["value"]
 	}
 	require.Equal(t, expectedSum, sum)
+
+	// Perform a query using the updated index in the schema.
+	q2 := `
+	{
+		names(func: anyofterms(name, Mike)) {
+			name
+		}
+	}`
+	resMap = make(map[string][]map[string]int)
+	resp, err = testutil.RetryQuery(dg, q2)
+	require.NoError(t, err)
+
+	// Trying to perform a query using the address index should not work since that
+	// predicate was deleted.
+	q3 := `
+	{
+		addresses(func: anyofterms(address, Mike)) {
+			address
+		}
+	}`
+	resMap = make(map[string][]map[string]int)
+	resp, err = testutil.RetryBadQuery(dg, q3)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "Attribute address is not indexed")
 }
 
 func waitForSnapshot(t *testing.T, prevSnapTs uint64) uint64 {