Function streamBadger

dgraph/cmd/dgraphimport/import_client.go:215–266

streamBadger runs a BadgerDB stream to send key-value pairs to the specified group. It creates a new stream at the maximum sequence number and sends the data to the specified group. It also sends a final 'done' signal to mark completion.

(ctx context.Context, ps *badger.DB, out api.Dgraph_StreamExtSnapshotClient, groupId uint32)

// It creates a new stream at the maximum sequence number and sends the data to the specified group.
// It also sends a final 'done' signal to mark completion.
func streamBadger(ctx context.Context, ps *badger.DB, out api.Dgraph_StreamExtSnapshotClient, groupId uint32) error {
	stream := ps.NewStreamAt(math.MaxUint64)
	stream.LogPrefix = "[import] Sending external snapshot to group [" + fmt.Sprintf("%d", groupId) + "]"
	stream.KeyToList = nil
	stream.Send = func(buf *z.Buffer) error {
		p := &api.StreamPacket{Data: buf.Bytes()}
		if err := out.Send(&api.StreamExtSnapshotRequest{Pkt: p}); err != nil && !errors.Is(err, io.EOF) {
			return fmt.Errorf("failed to send data chunk: %w", err)
		}
		if _, err := out.Recv(); err != nil {
			return fmt.Errorf("failed to receive response for group ID [%v] from the server: %w", groupId, err)
		}
		glog.Infof("[import] Group [%v]: Received ACK for sending data chunk", groupId)

		return nil
	}

	// Execute the stream process
	if err := stream.Orchestrate(ctx); err != nil {
		return fmt.Errorf("stream orchestration failed for group [%v]: %w, badger path: %s", groupId, err, ps.Opts().Dir)
	}

	// Send the final 'done' signal to mark completion
	glog.Infof("[import] Sending completion signal for group [%d]", groupId)
	done := &api.StreamPacket{Done: true}

	if err := out.Send(&api.StreamExtSnapshotRequest{Pkt: done}); err != nil && !errors.Is(err, io.EOF) {
		return fmt.Errorf("failed to send 'done' signal for group [%d]: %w", groupId, err)
	}

	for {
		if ctx.Err() != nil {
			return ctx.Err()
		}
		resp, err := out.Recv()
		if errors.Is(err, io.EOF) {
			return fmt.Errorf("server closed stream before Finish=true for group [%d]", groupId)
		}
		if err != nil {
			return fmt.Errorf("failed to receive final response for group ID [%v] from the server: %w", groupId, err)
		}
		if resp.Finish {
			glog.Infof("[import] Group [%v]: Received final Finish=true", groupId)
			break
		}
		glog.Infof("[import] Group [%v]: Waiting for Finish=true, got interim ACK", groupId)
	}

	glog.Infof("[import] Group [%v]: Received ACK for sending completion signal", groupId)

	return nil
}