Avoid race condition between predicate move and commit (#2392)

- The reason for bug #2338 was that there was a race condition between a mutation and predicate move. Zero was not checking if a predicate is under move before allowing a commit. Thus, a mutation could get proposed in a group, then a move starts, and get committed by Zero (after the move starts).
- This change this issue by ensuring that Zero checks if a predicate is being moved, before allowing commit.
- Any pending transactions are also cancelled once the move starts, so this would only happen as part of a race condition and not afterward.

Mechanism:
- Send the real keys back to Zero, as part of Transaction Context.
- Zero uses these keys to parse the predicate, and checks if that predicate is currently moving. If so, it would abort the transaction.
- Also, check for `_predicate_` being moved. For some reason, if we don't consider this predicate, we could still lose data.
- Before doing a mutation in Dgraph alpha, check if that tablet can be written to.
- Loop until all transactions corresponding to the predicate move are aborted. Only then start the move.

Tangential changes:
- Update the port number for bank integration test.
- Remove the separate key value or clean channel. Make it run as part of the main Node.Run loop.
- Add a max function.
- Small refactoring here and there.

contrib/integration/bank/main.go

@@ -175,7 +175,7 @@ func (s *State) loop(wg *sync.WaitGroup) {
func main() {
	flag.Parse()

	conn, err := grpc.Dial("localhost:9081", grpc.WithInsecure())
	conn, err := grpc.Dial("localhost:9080", grpc.WithInsecure())
	if err != nil {
		log.Fatal(err)
	}

dgraph/cmd/zero/oracle.go

@@ -8,6 +8,7 @@
package zero

import (
	"encoding/base64"
	"errors"
	"math/rand"
	"time"
@@ -284,6 +285,38 @@ func (s *Server) commit(ctx context.Context, src *api.TxnContext) error {
		return s.proposeTxn(ctx, src)
	}

	// Check if any of these tablets is being moved. If so, abort the transaction.
	preds := make(map[string]struct{})
	// _predicate_ would never be part of conflict detection, so keys corresponding to any
	// modifications to this predicate would not be sent to Zero. But, we still need to abort
	// transactions which are coming in, while this predicate is being moved. This means that if
	// _predicate_ expansion is enabled, and a move for this predicate is happening, NO transactions
	// across the entire cluster would commit. Sorry! But if we don't do this, we might lose commits
	// which sneaked in during the move.
	preds["_predicate_"] = struct{}{}

	for _, k := range src.Keys {
		key, err := base64.StdEncoding.DecodeString(k)
		if err != nil {
			continue
		}
		pk := x.Parse(key)
		if pk != nil {
			preds[pk.Attr] = struct{}{}
		}
	}
	for pred := range preds {
		tablet := s.ServingTablet(pred)
		if tablet == nil || tablet.GetReadOnly() {
			src.Aborted = true
			return s.proposeTxn(ctx, src)
		}
	}

	// TODO: We could take fingerprint of the keys, and store them in uint64, allowing the rowCommit
	// map to be keyed by uint64, which would be cheaper. But, unsure about the repurcussions of
	// that. It would save some memory. So, worth a try.

	var num intern.Num
	num.Val = 1
	assigned, err := s.lease(ctx, &num, true)

dgraph/cmd/zero/tablet.go

@@ -64,6 +64,10 @@ func (s *Server) rebalanceTablets() {
}

func (s *Server) movePredicate(predicate string, srcGroup, dstGroup uint32) error {
	// Typically move predicate is run only on leader. But, in this case, an external HTTP request
	// can also trigger a predicate move. We could render them invalid here by checking if this node
	// is actually the leader. But, I have noticed no side effects with allowing them to run, even
	// if this node is a follower node.
	tab := s.ServingTablet(predicate)
	x.AssertTruef(tab != nil, "Tablet to be moved: [%v] should not be nil", predicate)
	x.Printf("Going to move predicate: [%v], size: [%v] from group %d to %d\n", predicate,
@@ -81,6 +85,7 @@ func (s *Server) movePredicate(predicate string, srcGroup, dstGroup uint32) erro
				break
			}

			x.Printf("Sleeping before we run recovery for tablet move")
			// We might have initiated predicate move on some other node, give it some
			// time to get cancelled. On cancellation the other node would set the predicate
			// to write mode again and we need to be sure that it doesn't happen after we

posting/index.go

@@ -435,13 +435,11 @@ func deleteEntries(prefix []byte, remove func(key []byte) bool) error {
		if !remove(item.Key()) {
			continue
		}
		nkey := make([]byte, len(item.Key()))
		copy(nkey, item.Key())
		nkey := item.KeyCopy(nil)
		version := item.Version()

		txn := pstore.NewTransactionAt(version, true)
		// Purge doesn't delete anything, so write an empty pl
		txn.SetWithDiscard(nkey, nil, BitEmptyPosting)
		txn.Delete(nkey)
		wg.Add(1)
		err := txn.CommitAt(version, func(e error) {
			defer wg.Done()

posting/mvcc.go

@@ -10,9 +10,9 @@ package posting
import (
	"bytes"
	"context"
	"encoding/base64"
	"math"
	"sort"
	"strconv"
	"sync"
	"sync/atomic"
	"time"
@@ -21,7 +21,6 @@ import (
	"github.com/dgraph-io/dgo/protos/api"
	"github.com/dgraph-io/dgraph/protos/intern"
	"github.com/dgraph-io/dgraph/x"
	farm "github.com/dgryski/go-farm"
)

var (
@@ -216,8 +215,10 @@ func (t *Txn) Fill(ctx *api.TxnContext) {
	for i := t.nextKeyIdx; i < len(t.deltas); i++ {
		d := t.deltas[i]
		if d.checkConflict {
			fp := farm.Fingerprint64(d.key)
			ctx.Keys = append(ctx.Keys, strconv.FormatUint(fp, 36))
			// Instead of taking a fingerprint of the keys, send the whole key to Zero. So, Zero can
			// parse the key and check if that predicate is undergoing a move, hence avoiding #2338.
			k := base64.StdEncoding.EncodeToString(d.key)
			ctx.Keys = append(ctx.Keys, k)
		}
	}
	t.nextKeyIdx = len(t.deltas)

worker/draft.go

@@ -372,34 +372,7 @@ func (n *node) applyConfChange(e raftpb.Entry) {
	groups().triggerMembershipSync()
}

type KeyValueOrCleanProposal struct {
	raftIdx  uint64
	proposal *intern.Proposal
}

func (n *node) processKeyValueOrCleanProposals(
	kvChan chan KeyValueOrCleanProposal) {
	// Run KeyValueProposals and CleanPredicate one by one always.
	// During predicate move we first clean the predicate and then
	// propose key values, we wait for clean predicate to be done before
	// we propose key values. But during replay if we run these proposals
	// in goroutine then we will have no such guarantees so always run
	// them sequentially.
	for e := range kvChan {
		if len(e.proposal.Kv) > 0 {
			n.processKeyValues(e.raftIdx, e.proposal.Key, e.proposal.Kv)
		} else if len(e.proposal.CleanPredicate) > 0 {
			n.deletePredicate(e.raftIdx, e.proposal.Key, e.proposal.CleanPredicate)
		} else {
			x.Fatalf("Unknown proposal, %+v\n", e.proposal)
		}
	}
}

func (n *node) processApplyCh() {
	kvChan := make(chan KeyValueOrCleanProposal, 1000)
	go n.processKeyValueOrCleanProposals(kvChan)

	for e := range n.applyCh {
		if len(e.Data) == 0 {
			// This is not in the proposal map
@@ -443,30 +416,27 @@ func (n *node) processApplyCh() {
		if proposal.Mutations != nil {
			// syncmarks for this shouldn't be marked done until it's comitted.
			n.sch.schedule(proposal, e.Index)

		} else if len(proposal.Kv) > 0 {
			kvChan <- KeyValueOrCleanProposal{
				raftIdx:  e.Index,
				proposal: proposal,
			}
			n.processKeyValues(e.Index, proposal.Key, proposal.Kv)

		} else if proposal.State != nil {
			// This state needn't be snapshotted in this group, on restart we would fetch
			// a state which is latest or equal to this.
			groups().applyState(proposal.State)
			// When proposal is done it emits done watermarks.
			posting.TxnMarks().Done(e.Index)
			n.props.Done(proposal.Key, nil)

		} else if len(proposal.CleanPredicate) > 0 {
			kvChan <- KeyValueOrCleanProposal{
				raftIdx:  e.Index,
				proposal: proposal,
			}
			n.deletePredicate(e.Index, proposal.Key, proposal.CleanPredicate)

		} else if proposal.TxnContext != nil {
			go n.commitOrAbort(e.Index, proposal.Key, proposal.TxnContext)
		} else {
			x.Fatalf("Unknown proposal")
		}
	}
	close(kvChan)
}

func (n *node) commitOrAbort(index uint64, pid string, tctx *api.TxnContext) {
@@ -492,11 +462,11 @@ func (n *node) deletePredicate(index uint64, pid string, predicate string) {
	n.props.Done(pid, err)
}

func (n *node) processKeyValues(index uint64, pid string, kvs []*intern.KV) error {
	ctx, _ := n.props.CtxAndTxn(pid)
func (n *node) processKeyValues(index uint64, pkey string, kvs []*intern.KV) error {
	ctx, _ := n.props.CtxAndTxn(pkey)
	err := populateKeyValues(ctx, kvs)
	posting.TxnMarks().Done(index)
	n.props.Done(pid, err)
	n.props.Done(pkey, err)
	return nil
}

@@ -865,7 +835,7 @@ func (n *node) joinPeers() error {

	gconn := pl.Get()
	c := intern.NewRaftClient(gconn)
	x.Printf("Calling JoinCluster")
	x.Printf("Calling JoinCluster via leader: %s", pl.Addr)
	if _, err := c.JoinCluster(n.ctx, n.RaftContext); err != nil {
		return x.Errorf("Error while joining cluster: %+v\n", err)
	}
@@ -935,11 +905,11 @@ func (n *node) InitAndStartNode(wal *raftwal.Wal) {
		n.SetRaft(raft.RestartNode(n.Cfg))
	} else {
		x.Printf("New Node for group: %d\n", n.gid)
		if _, hasPeer := groups().MyPeer(); hasPeer {
		if peerId, hasPeer := groups().MyPeer(); hasPeer {
			// Get snapshot before joining peers as it can take time to retrieve it and we dont
			// want the quorum to be inactive when it happens.

			x.Println("Retrieving snapshot.")
			x.Println("Retrieving snapshot from peer: %d", peerId)
			n.retryUntilSuccess(n.retrieveSnapshot, time.Second)

			x.Println("Trying to join peers.")

worker/groups.go

@@ -306,6 +306,14 @@ func (g *groupi) BelongsTo(key string) uint32 {
	return 0
}

func (g *groupi) ServesTabletRW(key string) bool {
	tablet := g.Tablet(key)
	if tablet != nil && !tablet.ReadOnly && tablet.GroupId == groups().groupId() {
		return true
	}
	return false
}

func (g *groupi) ServesTablet(key string) bool {
	tablet := g.Tablet(key)
	if tablet != nil && tablet.GroupId == groups().groupId() {

worker/mutation.go

@@ -38,13 +38,12 @@ func deletePredicateEdge(edge *intern.DirectedEdge) bool {
	return edge.Entity == 0 && bytes.Equal(edge.Value, []byte(x.Star))
}

// runMutation goes through all the edges and applies them. It returns the
// mutations which were not applied in left.
// runMutation goes through all the edges and applies them.
func runMutation(ctx context.Context, edge *intern.DirectedEdge, txn *posting.Txn) error {
	if tr, ok := trace.FromContext(ctx); ok {
		tr.LazyPrintf("In run mutations")
	}
	if !groups().ServesTablet(edge.Attr) {
	if !groups().ServesTabletRW(edge.Attr) {
		// Don't assert, can happen during replay of raft logs if server crashes immediately
		// after predicate move and before snapshot.
		return errUnservedTablet
@@ -411,10 +410,8 @@ func fillTxnContext(tctx *api.TxnContext, gid uint32, startTs uint64) {
	// applied watermark can be less than this proposal's index so return the maximum.
	// For some proposals like dropPredicate, we don't store them in txns map, so we
	// don't know the raft index. For them we would return applied watermark.
	if x := node.Applied.DoneUntil(); x > index {
		index = x
	}
	tctx.LinRead.Ids[gid] = index
	doneUntil := node.Applied.DoneUntil()
	tctx.LinRead.Ids[gid] = x.Max(index, doneUntil)
}

// proposeOrSend either proposes the mutation if the node serves the group gid or sends it to