Use rampMeter for Executor

worker/draft.go

@@ -852,17 +852,20 @@ func (n *node) proposeSnapshot(discardN int) error {
 	return n.Raft().Propose(n.ctx, data)
 }
 
-const maxPendingSize int64 = 64 << 20 // in bytes.
+const (
+	maxPendingSize int64 = 64 << 20 // in bytes.
+	nodeApplyChan        = "raft node applyCh"
+)
 
-func (n *node) rampMeter() {
+func rampMeter(address *int64, maxSize int64, component string) {
 	start := time.Now()
 	defer func() {
 		if dur := time.Since(start); dur > time.Second {
-			glog.Infof("Blocked pushing to applyCh for %v", dur.Round(time.Millisecond))
+			glog.Infof("Blocked pushing to %s for %v", component, dur.Round(time.Millisecond))
 		}
 	}()
 	for {
-		if atomic.LoadInt64(&n.pendingSize) <= maxPendingSize {
+		if atomic.LoadInt64(address) <= maxSize {
 			return
 		}
 		time.Sleep(3 * time.Millisecond)
@@ -1171,7 +1174,7 @@ func (n *node) Run() {
 				// Apply the meter this before adding size to pending size so some crazy big
 				// proposal can be pushed to applyCh. If this do this after adding its size to
 				// pending size, we could block forever in rampMeter.
-				n.rampMeter()
+				rampMeter(&n.pendingSize, maxPendingSize, nodeApplyChan)
 				var pendingSize int64
 				for _, p := range proposals {
 					pendingSize += int64(p.Size())

worker/executor.go

@@ -21,6 +21,7 @@ package worker
 import (
 	"context"
 	"sync"
+	"sync/atomic"
 
 	"github.com/dgraph-io/badger/v2/y"
 	"github.com/dgraph-io/dgraph/posting"
@@ -35,6 +36,8 @@ type subMutation struct {
 }
 
 type executor struct {
+	pendingSize int64
+
 	sync.RWMutex
 	predChan map[string]chan *subMutation
 	closer   *y.Closer
@@ -54,8 +57,10 @@ func (e *executor) processMutationCh(ch chan *subMutation) {
 
 	writer := posting.NewTxnWriter(pstore)
 	for payload := range ch {
+		var esize int64
 		ptxn := posting.NewTxn(payload.startTs)
 		for _, edge := range payload.edges {
+			esize += int64(edge.Size())
 			for {
 				err := runMutation(payload.ctx, edge, ptxn)
 				if err == nil {
@@ -74,6 +79,8 @@ func (e *executor) processMutationCh(ch chan *subMutation) {
 		if err := writer.Wait(); err != nil {
 			glog.Errorf("Error while waiting for writes: %v", err)
 		}
+
+		atomic.AddInt64(&e.pendingSize, -esize)
 	}
 }
 
@@ -99,9 +106,16 @@ func (e *executor) getChannelUnderLock(pred string) (ch chan *subMutation) {
 	return ch
 }
 
+const (
+	maxPendingEdgesSize int64 = 64 << 20
+	executorAddEdges          = "executor.addEdges"
+)
+
 func (e *executor) addEdges(ctx context.Context, startTs uint64, edges []*pb.DirectedEdge) {
-	payloadMap := make(map[string]*subMutation)
+	rampMeter(&e.pendingSize, maxPendingEdgesSize, executorAddEdges)
 
+	payloadMap := make(map[string]*subMutation)
+	var esize int64
 	for _, edge := range edges {
 		payload, ok := payloadMap[edge.Attr]
 		if !ok {
@@ -112,6 +126,7 @@ func (e *executor) addEdges(ctx context.Context, startTs uint64, edges []*pb.Dir
 			payload = payloadMap[edge.Attr]
 		}
 		payload.edges = append(payload.edges, edge)
+		esize += int64(edge.Size())
 	}
 
 	// Lock() in case the channel gets closed from underneath us.
@@ -127,4 +142,5 @@ func (e *executor) addEdges(ctx context.Context, startTs uint64, edges []*pb.Dir
 		}
 	}
 
+	atomic.AddInt64(&e.pendingSize, esize)
 }