Repository: dgraph-io/dgraph
File: worker/import.go (Go)

Code before:
	eg, errGCtx := errgroup.WithContext(stream.Context())
	for _, member := range groups().state.Groups[node.gid].Members {
		if member.Addr == node.MyAddr {
			eg.Go(func() error {
				if err := ps.runLocalSubscriber(errGCtx, stream); err != nil {
					glog.Errorf("[import:flush] failed to run local subscriber: %v", err)
					updateNodeStatus(&ps.RWMutex, successfulNodes, member.Addr, false)
					return err
				}

				updateNodeStatus(&ps.RWMutex, successfulNodes, member.Addr, true)
				return nil
			})
			continue
		}

		if forward {
			// We are not going to return any error from here because we care about the majority of nodes.
			// If the majority of nodes are able to receive the data, the remaining ones can catch up later.
			glog.Infof("[import] Streaming external snapshot to [%v] from [%v] forward [%v]", member.Addr, node.MyAddr)
			eg.Go(func() error {
				glog.Infof(`[import:forward] streaming external snapshot to [%v] from [%v]`, member.Addr, node.MyAddr)
				if member.AmDead {
					glog.Infof(`[import:forward] [%v] is dead, skipping`, member.Addr)
					return nil
				}

				pl, err := conn.GetPools().Get(member.Addr)
				if err != nil {
					updateNodeStatus(&ps.RWMutex, successfulNodes, member.Addr, false)
					glog.Errorf("connection error to [%v]: %v", member.Addr, err)
					return nil
				}

				c := pb.NewWorkerClient(pl.Get())
				peerStream, err := c.StreamExtSnapshot(errGCtx)
				if err != nil {
					updateNodeStatus(&ps.RWMutex, successfulNodes, member.Addr, false)
					glog.Errorf("failed to establish stream with peer %v: %v", member.Addr, err)
					return nil
				}
				defer func() {
					if err := peerStream.CloseSend(); err != nil {
						glog.Errorf("[import:forward] failed to close stream with peer [%v]: %v", member.Addr, err)
					}
				}()

				forwardReq := &apiv2.StreamExtSnapshotRequest{Forward: false}
				if err := peerStream.Send(forwardReq); err != nil {
					updateNodeStatus(&ps.RWMutex, successfulNodes, member.Addr, false)
					glog.Errorf("failed to send forward request: %v", err)

Review comment (suggestion):
[nitpick] The log message should be more precise. Consider changing 'trying to update' to 'updating' to match the action being performed.
```suggestion
	glog.Info("[import] updating the import mode to false")
```

Code after:

	}
}

func (w *grpcWorker) UpdateExtSnapshotStreamingState(ctx context.Context,
	req *api.UpdateExtSnapshotStreamingStateRequest) (*pb.Status, error) {
	if req == nil {
		return nil, errors.New("UpdateExtSnapshotStreamingStateRequest must not be nil")
	}

	if req.Start && req.Finish {
		return nil, errors.New("UpdateExtSnapshotStreamingStateRequest cannot have both Start and Finish set to true")
	}

	glog.Infof("[import] Applying import mode proposal: %+v", req)
	err := groups().Node.proposeAndWait(ctx, &pb.Proposal{ExtSnapshotState: req})

	return &pb.Status{}, err
}

// StreamExtSnapshot handles the stream of key-value pairs sent from proxy alpha.
// It receives a Forward flag from the stream to determine if the current node is the leader.
// If the node is the leader (Forward is true), it streams the data to its followers.
// Otherwise, it simply writes the data to BadgerDB and flushes it.
func (w *grpcWorker) StreamExtSnapshot(stream pb.Worker_StreamExtSnapshotServer) error {
	glog.Info("[import] trying to update the import mode to false")
	defer x.ExtSnapshotStreamingState(false)

	// Receive the first message to check the Forward flag.
	// If Forward is true, this node is the leader and should forward the stream to its followers.
	// If Forward is false, the node just writes and flushes the data.
	forwardReq, err := stream.Recv()
	if err != nil {
		return err
	}

	glog.Infof("[import] received forward flag: %v", forwardReq.Forward)
	return streamInGroup(stream, forwardReq.Forward)
}

// postStreamProcessing handles the post-stream processing of data received from the buffer into the local BadgerDB.
// It loads the schema, updates the membership state, informs zero about tablets, resets caches, applies initial schema,
// applies initial types, and resets the GQL schema store.
func postStreamProcessing(ctx context.Context) error {
	glog.Info("[import:flush] post stream processing")
	if err := schema.LoadFromDb(ctx); err != nil {
		return errors.Wrapf(err, "cannot load schema after streaming data")
	}
	if err := UpdateMembershipState(ctx); err != nil {
		return errors.Wrapf(err, "cannot update membership state after streaming data")
	}