Method ReceivePredicate

worker/predicate_move.go:120–187

Returns count which can be used to verify whether we have moved all keys for a predicate or not.

(stream pb.Worker_ReceivePredicateServer)

// Returns count which can be used to verify whether we have moved all keys
// for a predicate or not.
func (w *grpcWorker) ReceivePredicate(stream pb.Worker_ReceivePredicateServer) error {
	if !groups().Node.AmLeader() {
		return errors.Errorf("ReceivePredicate failed: Not the leader of group")
	}
	// No new deletion/background cleanup would start after we start streaming tablet,
	// so all the proposals for a particular tablet would atmost wait for deletion of
	// single tablet. Only leader needs to do this.
	mu := groups().blockDeletes
	mu.Lock()
	defer mu.Unlock()

	// Values can be pretty big so having less buffer is safer.
	kvs := make(chan *pb.KVS, 3)
	che := make(chan error, 1)
	// We can use count to check the number of posting lists returned in tests.
	count := 0
	ctx := stream.Context()
	payload := &api.Payload{}

	glog.Infof("Got ReceivePredicate. Group: %d. Am leader: %v",
		groups().groupId(), groups().Node.AmLeader())

	go func() {
		// Takes care of throttling and batching.
		che <- batchAndProposeKeyValues(ctx, kvs)
	}()
	for {
		kvBuf, err := stream.Recv()
		if err == io.EOF {
			payload.Data = []byte(fmt.Sprintf("%d", count))
			if err := stream.SendAndClose(payload); err != nil {
				glog.Errorf("Received %d keys. Error in loop: %v", count, err)
				return err
			}
			break
		}
		if err != nil {
			glog.Errorf("Received %d keys. Error in loop: %v", count, err)
			return err
		}
		glog.V(2).Infof("Received batch of size: %s", humanize.IBytes(uint64(len(kvBuf.Data))))

		buf := z.NewBufferSlice(kvBuf.Data)
		if err := buf.SliceIterate(func(_ []byte) error {
			count++
			return nil
		}); err != nil {
			glog.Errorf("error while counting in buf: %v\n", err)
			return err
		}

		select {
		case kvs <- kvBuf:
		case <-ctx.Done():
			close(kvs)
			<-che
			glog.Infof("Received %d keys. Context deadline\n", count)
			return ctx.Err()