Add support for HA and multishard functionality in import APIs

Author: shivaji-kharse

dgraph/cmd/dgraphimport/import_client.go

@@ -7,16 +7,13 @@ package dgraphimport
 
 import (
 	"context"
-	"errors"
 	"fmt"
-	"io"
-	"math"
 	"os"
 	"path/filepath"
 
 	"github.com/dgraph-io/badger/v4"
 	apiv2 "github.com/dgraph-io/dgo/v250/protos/api.v2"
-	"github.com/dgraph-io/ristretto/v2/z"
+	"github.com/hypermodeinc/dgraph/v25/worker"
 
 	"github.com/golang/glog"
 	"golang.org/x/sync/errgroup"
@@ -48,10 +45,13 @@ func Import(ctx context.Context, endpoint string, opts grpc.DialOption, bulkOutD
 }
 
 // startPDirStream initiates a snapshot stream session with the Dgraph server.
-func startPDirStream(ctx context.Context, dc apiv2.DgraphClient) (*apiv2.InitiatePDirStreamResponse, error) {
+func startPDirStream(ctx context.Context, dc apiv2.DgraphClient) (*apiv2.DrainModeResponse, error) {
 	glog.Info("Initiating pdir stream")
-	req := &apiv2.InitiatePDirStreamRequest{}
-	resp, err := dc.InitiatePDirStream(ctx, req)
+	req := &apiv2.DrainModeRequest{
+		DrainMode: true,
+		DropData:  false,
+	}
+	resp, err := dc.UpdateClusterDrainMode(ctx, req)
 	if err != nil {
 		glog.Errorf("failed to initiate pdir stream: %v", err)
 		return nil, fmt.Errorf("failed to initiate pdir stream: %v", err)
@@ -63,7 +63,7 @@ func startPDirStream(ctx context.Context, dc apiv2.DgraphClient) (*apiv2.Initiat
 // sendPDir takes a p directory and a set of group IDs and streams the data from the
 // p directory to the corresponding group IDs. It first scans the provided directory for
 // subdirectories named with numeric group IDs.
-func sendPDir(ctx context.Context, dg apiv2.DgraphClient, baseDir string, groups []uint32) error {
+func sendPDir(ctx context.Context, dc apiv2.DgraphClient, baseDir string, groups []uint32) error {
 	glog.Infof("Starting to stream pdir from directory: %s", baseDir)
 
 	errG, ctx := errgroup.WithContext(ctx)
@@ -74,20 +74,39 @@ func sendPDir(ctx context.Context, dg apiv2.DgraphClient, baseDir string, groups
 			if _, err := os.Stat(pDir); err != nil {
 				return fmt.Errorf("p directory does not exist for group [%d]: [%s]", group, pDir)
 			}
-
 			glog.Infof("Streaming data for group [%d] from directory: [%s]", group, pDir)
-			if err := streamData(ctx, dg, pDir, group); err != nil {
+			if err := streamData(ctx, dc, pDir, group); err != nil {
 				glog.Errorf("Failed to stream data for groups [%v] from directory: [%s]: %v", group, pDir, err)
 				return err
 			}
 
 			return nil
 		})
 	}
-	if err := errG.Wait(); err != nil {
-		return err
+	if err1 := errG.Wait(); err1 != nil {
+		// If the p directory doesn't exist for this group, it indicates that
+		// streaming might be in progress to other groups. We disable drain mode
+		// to prevent interference and drop any streamed data to ensure a clean state.
+		req := &apiv2.DrainModeRequest{
+			DrainMode: false,
+			DropData:  true,
+		}
+		if _, err := dc.UpdateClusterDrainMode(context.Background(), req); err != nil {
+			return fmt.Errorf("failed to stream data :%v failed to off drain mode: %v", err1, err)
+		}
+
+		glog.Info("success fully disabled drain mode")
+		return err1
 	}
 
+	glog.Info("Completed streaming all pdirs")
+	req := &apiv2.DrainModeRequest{
+		DrainMode: false,
+		DropData:  false,
+	}
+	if _, err := dc.UpdateClusterDrainMode(context.Background(), req); err != nil {
+		return fmt.Errorf("failed to off drain mode: %v", err)
+	}
 	glog.Infof("Completed streaming all pdirs")
 	return nil
 }
@@ -125,34 +144,16 @@ func streamData(ctx context.Context, dg apiv2.DgraphClient, pdir string, groupId
 
 	// Configure and start the BadgerDB stream
 	glog.Infof("Starting BadgerDB stream for group [%d]", groupId)
-	stream := ps.NewStreamAt(math.MaxUint64)
-	stream.LogPrefix = fmt.Sprintf("Sending P dir for group [%d]", groupId)
-	stream.KeyToList = nil
-	stream.Send = func(buf *z.Buffer) error {
-		p := &apiv2.StreamPacket{Data: buf.Bytes()}
-		if err := out.Send(&apiv2.StreamPDirRequest{StreamPacket: p}); err != nil && !errors.Is(err, io.EOF) {
-			return fmt.Errorf("failed to send data chunk: %w", err)
-		}
-		return nil
-	}
+	// if err := RunBadgerStream(ctx, ps, out, groupId); err != nil {
+	// 	return fmt.Errorf("badger stream failed for group [%d]: %w", groupId, err)
+	// }
 
-	// Execute the stream process
-	if err := stream.Orchestrate(ctx); err != nil {
-		return fmt.Errorf("stream orchestration failed for group [%d]: %w", groupId, err)
+	if err := worker.RunBadgerStream(ctx, ps, out, groupId); err != nil {
+		return fmt.Errorf("badger stream failed for group [%d]: %w", groupId, err)
 	}
-
-	// Send the final 'done' signal to mark completion
-	glog.Infof("Sending completion signal for group [%d]", groupId)
-	done := &apiv2.StreamPacket{Done: true}
-
-	if err := out.Send(&apiv2.StreamPDirRequest{StreamPacket: done}); err != nil && !errors.Is(err, io.EOF) {
-		return fmt.Errorf("failed to send 'done' signal for group [%d]: %w", groupId, err)
-	}
-	// Wait for acknowledgment from the server
 	if _, err := out.CloseAndRecv(); err != nil {
 		return fmt.Errorf("failed to receive ACK for group [%d]: %w", groupId, err)
 	}
 	glog.Infof("Group [%d]: Received ACK ", groupId)
-
 	return nil
 }