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
 }

dgraph/cmd/dgraphimport/import_test.go

@@ -9,7 +9,10 @@ package dgraphimport
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"path/filepath"
+	"strconv"
+	"strings"
 	"testing"
 
 	"github.com/hypermodeinc/dgraph/v25/dgraphapi"
@@ -91,39 +94,223 @@ func TestDrainModeAfterStartSnapshotStream(t *testing.T) {
 // TestImportApis tests import functionality with different cluster configurations
 func TestImportApis(t *testing.T) {
 	tests := []struct {
-		name           string
-		bulkAlphas     int
-		targetAlphas   int
-		replicasFactor int
+		name             string
+		bulkAlphas       int  // Number of alphas in source cluster
+		targetAlphas     int  // Number of alphas in target cluster
+		replicasFactor   int  // Number of replicas for each group
+		downAlphas       int  // Number of alphas to be shutdown
+		negativeTestCase bool // True if this is an expected failure case
+		description      string
+		err              string
 	}{
-		{"SingleGroupSingleAlpha", 1, 1, 1},
-		{"TwoGroupsSingleAlpha", 2, 2, 1},
-		{"ThreeGroupsSingleAlpha", 3, 3, 1},
+		{
+			name:             "SingleGroupShutTwoAlphasPerGroup",
+			bulkAlphas:       1,
+			targetAlphas:     3,
+			replicasFactor:   3,
+			downAlphas:       2,
+			negativeTestCase: true,
+			description:      "Single group with 3 alphas, shutdown 2 alphas",
+			err:              "failed to initiate pdir stream",
+		},
+		{
+			name:             "TwoGroupShutTwoAlphasPerGroup",
+			bulkAlphas:       2,
+			targetAlphas:     6,
+			replicasFactor:   3,
+			downAlphas:       2,
+			negativeTestCase: true,
+			description:      "Two groups with 3 alphas each, shutdown 2 alphas per group",
+			err:              "failed to initiate pdir stream",
+		},
+		{
+			name:             "TwoGroupShutTwoAlphasPerGroupNoPDir",
+			bulkAlphas:       1,
+			targetAlphas:     6,
+			replicasFactor:   3,
+			downAlphas:       0,
+			negativeTestCase: true,
+			description:      "Two groups with 3 alphas each, 1 p directory is not available",
+			err:              "p directory does not exist for group [2]",
+		},
+		{
+			name:             "ThreeGroupShutTwoAlphasPerGroup",
+			bulkAlphas:       3,
+			targetAlphas:     9,
+			replicasFactor:   3,
+			downAlphas:       2,
+			negativeTestCase: true,
+			description:      "Three groups with 3 alphas each, shutdown 2 alphas per group",
+			err:              "failed to initiate pdir stream",
+		},
+		{
+			name:             "SingleGroupShutOneAlpha",
+			bulkAlphas:       1,
+			targetAlphas:     3,
+			replicasFactor:   3,
+			downAlphas:       1,
+			negativeTestCase: false,
+			description:      "Single group with multiple alphas, shutdown 1 alpha",
+			err:              "",
+		},
+		{
+			name:             "TwoGroupShutOneAlphaPerGroup",
+			bulkAlphas:       2,
+			targetAlphas:     6,
+			replicasFactor:   3,
+			downAlphas:       1,
+			negativeTestCase: false,
+			description:      "Multiple groups with multiple alphas, shutdown 2 alphas per group",
+			err:              "",
+		},
+		{
+			name:             "ThreeGroupShutOneAlphaPerGroup",
+			bulkAlphas:       3,
+			targetAlphas:     9,
+			replicasFactor:   3,
+			downAlphas:       1,
+			negativeTestCase: false,
+			description:      "Three groups with 3 alphas each, shutdown 1 alpha per group",
+			err:              "",
+		},
+		{
+			name:             "SingleGroupAllAlphasOnline",
+			bulkAlphas:       1,
+			targetAlphas:     3,
+			replicasFactor:   3,
+			downAlphas:       0,
+			negativeTestCase: false,
+			description:      "Single group with multiple alphas, all alphas are online",
+			err:              "",
+		},
+		{
+			name:             "TwoGroupAllAlphasOnline",
+			bulkAlphas:       2,
+			targetAlphas:     6,
+			replicasFactor:   3,
+			downAlphas:       0,
+			negativeTestCase: false,
+			description:      "Multiple groups with multiple alphas, all alphas are online",
+			err:              "",
+		},
+		{
+			name:             "ThreeGroupAllAlphasOnline",
+			bulkAlphas:       3,
+			targetAlphas:     9,
+			replicasFactor:   3,
+			downAlphas:       0,
+			negativeTestCase: false,
+			description:      "Three groups with 3 alphas each, all alphas are online",
+			err:              "",
+		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			runImportTest(t, tt.bulkAlphas, tt.targetAlphas, tt.replicasFactor)
+			if tt.negativeTestCase {
+				t.Logf("Running negative test case: %s", tt.description)
+			} else {
+				t.Logf("Running test case: %s", tt.description)
+			}
+			runImportTest(t, tt.bulkAlphas, tt.targetAlphas, tt.replicasFactor, tt.downAlphas, tt.negativeTestCase, tt.err)
 		})
 	}
 }
 
-func runImportTest(t *testing.T, bulkAlphas, targetAlphas, replicasFactor int) {
+func runImportTest(t *testing.T, bulkAlphas, targetAlphas, replicasFactor, numDownAlphas int, negative bool, errStr string) {
 	bulkCluster, baseDir := setupBulkCluster(t, bulkAlphas)
 	defer func() { bulkCluster.Cleanup(t.Failed()) }()
 
 	targetCluster, gc, gcCleanup := setupTargetCluster(t, targetAlphas, replicasFactor)
 	defer func() { targetCluster.Cleanup(t.Failed()) }()
 	defer gcCleanup()
 
+	_, err := gc.Query("schema{}")
+	require.NoError(t, err)
+
 	url, err := targetCluster.GetAlphaGrpcEndpoint(0)
 	require.NoError(t, err)
 	outDir := filepath.Join(baseDir, "out")
 
+	// Get group information for all alphas
+	_, cleanup, err := targetCluster.Client()
+	require.NoError(t, err)
+	defer cleanup()
+
+	// Get health status for all instances
+	hc, err := targetCluster.HTTPClient()
+	require.NoError(t, err)
+
+	healthResp, err := hc.HealthForInstance()
+	require.NoError(t, err)
+
+	// Parse health response to get group information
+	var guardianResp struct {
+		Health []struct {
+			Instance string
+			Address  string
+			LastEcho int64
+			Status   string
+			Version  string
+			UpTime   int64
+			Group    string
+		}
+	}
+	require.NoError(t, json.Unmarshal(healthResp, &guardianResp))
+
+	// Group alphas by their group number
+	alphaGroups := make(map[string][]int)
+	for _, h := range guardianResp.Health {
+		if strings.Contains(h.Instance, "zero") || strings.Contains(h.Address, "alpha0") {
+			continue
+		}
+		// Extract alpha number from address format like "alpha2:7080"
+		alphaNum := strings.TrimPrefix(h.Address, "alpha")
+		alphaNum = strings.TrimSuffix(alphaNum, ":7080")
+
+		alphaID, err := strconv.Atoi(alphaNum)
+		require.NoError(t, err)
+		alphaGroups[h.Group] = append(alphaGroups[h.Group], alphaID)
+	}
+
+	// Shutdown specified number of alphas from each group
+	for group, alphas := range alphaGroups {
+		for i := 0; i < numDownAlphas; i++ {
+			alphaID := alphas[i]
+			t.Logf("Shutting down alpha %d from group %s", alphaID, group)
+			require.NoError(t, targetCluster.StopAlpha(alphaID))
+		}
+	}
+
+	if negative {
+		err := Import(context.Background(), url, grpc.WithTransportCredentials(insecure.NewCredentials()), outDir)
+		require.Error(t, err)
+		fmt.Println("Error: ", err)
+		require.ErrorContains(t, err, errStr)
+		return
+	}
+
 	require.NoError(t, Import(context.Background(), url,
 		grpc.WithTransportCredentials(insecure.NewCredentials()), outDir))
 
-	verifyImportResults(t, gc)
+	for group, alphas := range alphaGroups {
+		for i := 0; i < numDownAlphas; i++ {
+			alphaID := alphas[i]
+			t.Logf("Shutting down alpha %v from group %v", alphaID, group)
+			require.NoError(t, targetCluster.StartAlpha(alphaID))
+		}
+	}
+
+	require.NoError(t, targetCluster.HealthCheck(false))
+
+	fmt.Println("Import completed")
+
+	for i := 0; i < targetAlphas; i++ {
+		gc, cleanup, err := targetCluster.AlphaClient(i)
+		require.NoError(t, err)
+		defer cleanup()
+		verifyImportResults(t, gc)
+	}
 }
 
 // setupBulkCluster creates and configures a cluster for bulk loading data
@@ -156,7 +343,7 @@ func setupBulkCluster(t *testing.T, numAlphas int) (*dgraphtest.LocalCluster, st
 func setupTargetCluster(t *testing.T, numAlphas, replicasFactor int) (*dgraphtest.LocalCluster, *dgraphapi.GrpcClient, func()) {
 	conf := dgraphtest.NewClusterConfig().
 		WithNumAlphas(numAlphas).
-		WithNumZeros(1).
+		WithNumZeros(3).
 		WithReplicas(replicasFactor)
 
 	cluster, err := dgraphtest.NewLocalCluster(conf)