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"
 	apiv25 "github.com/dgraph-io/dgo/v250/protos/api.v25"
-	"github.com/dgraph-io/ristretto/v2/z"
+	"github.com/hypermodeinc/dgraph/v25/worker"
 
 	"github.com/golang/glog"
 	"golang.org/x/sync/errgroup"
@@ -125,34 +122,16 @@ func streamData(ctx context.Context, dg apiv25.DgraphClient, pdir string, groupI
 
 	// 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 := &apiv25.StreamPacket{Data: buf.Bytes()}
-		if err := out.Send(&apiv25.StreamPDirRequest{StreamPacket: p}); err != nil && !errors.Is(err, io.EOF) {
-			return fmt.Errorf("failed to send data chunk: %w", err)
-		}
-		return nil
-	}
-
-	// Execute the stream process
-	if err := stream.Orchestrate(ctx); err != nil {
-		return fmt.Errorf("stream orchestration 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 := &apiv25.StreamPacket{Done: true}
+	// if err := RunBadgerStream(ctx, ps, out, groupId); err != nil {
+	// 	return fmt.Errorf("badger stream failed for group [%d]: %w", groupId, err)
+	// }
 
-	if err := out.Send(&apiv25.StreamPDirRequest{StreamPacket: done}); err != nil && !errors.Is(err, io.EOF) {
-		return fmt.Errorf("failed to send 'done' signal 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)
 	}
-	// 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

@@ -11,6 +11,7 @@ import (
 	"encoding/json"
 	"path/filepath"
 	"testing"
+	"time"
 
 	"github.com/hypermodeinc/dgraph/v25/dgraphapi"
 	"github.com/hypermodeinc/dgraph/v25/dgraphtest"
@@ -96,9 +97,9 @@ func TestImportApis(t *testing.T) {
 		targetAlphas   int
 		replicasFactor int
 	}{
-		{"SingleGroupSingleAlpha", 1, 1, 1},
-		{"TwoGroupsSingleAlpha", 2, 2, 1},
-		{"ThreeGroupsSingleAlpha", 3, 3, 1},
+		{"SingleGroupSingleAlpha", 1, 3, 3},
+		{"TwoGroupsSingleAlpha", 2, 6, 3},
+		{"ThreeGroupsSingleAlpha", 3, 9, 3},
 	}
 
 	for _, tt := range tests {
@@ -123,6 +124,7 @@ func runImportTest(t *testing.T, bulkAlphas, targetAlphas, replicasFactor int) {
 	require.NoError(t, Import(context.Background(), url,
 		grpc.WithTransportCredentials(insecure.NewCredentials()), outDir))
 
+	time.Sleep(time.Minute * 3)
 	verifyImportResults(t, gc)
 }
 

edgraph/server.go

@@ -1829,11 +1829,6 @@ func (s *ServerV25) StreamPDir(stream apiv25.Dgraph_StreamPDirServer) error {
 		return err
 	}
 
-	drainMode := &pb.DrainModeRequest{State: false}
-	if _, err := worker.ProposeDrain(stream.Context(), drainMode); err != nil {
-		return err
-	}
-
 	return nil
 }
 

protos/pb.proto

@@ -320,6 +320,7 @@ message Proposal {
   // Skipping 15 as it is used for uint64 key in master and might be needed later here.
   uint64 start_ts = 16;
   DrainModeRequest drainmode = 17;
+  ReqPDirStreamRequest reqpdirstream = 18;
 }
 
 message CDCState {
@@ -598,6 +599,7 @@ service Worker {
   rpc TaskStatus(TaskStatusRequest) returns (TaskStatusResponse) {}
   rpc ApplyDrainmode(DrainModeRequest) returns (Status) {}
   rpc InternalStreamPDir(stream api.v25.StreamPDirRequest) returns (api.v25.StreamPDirResponse) {}
+  rpc ReqPDirStream(ReqPDirStreamRequest) returns (stream api.v25.StreamPDirRequest) {}
 }
 
 message DrainModeRequest {
@@ -795,4 +797,8 @@ message TaskStatusResponse {
   uint64 task_meta = 1;
 }
 
+message ReqPDirStreamRequest {
+  string addr=2 ;
+}
+
 // vim: expandtab sw=2 ts=2

protos/pb/pb.pb.go

@@ -86,6 +86,8 @@ func (id op) String() string {
 		return "opBackup"
 	case opPredMove:
 		return "opPredMove"
+	case opStreamPDir:
+		return "opStreamPDir"
 	default:
 		return "opUnknown"
 	}
@@ -98,6 +100,7 @@ const (
 	opRestore
 	opBackup
 	opPredMove
+	opStreamPDir
 )
 
 // startTask is used for the tasks that do not require tracking of timestamp.
@@ -693,6 +696,44 @@ func (n *node) applyCommitted(proposal *pb.Proposal, key uint64) error {
 		x.UpdateDrainingMode(proposal.Drainmode.State)
 		return nil
 
+	case proposal.Reqpdirstream != nil:
+		defer x.UpdateDrainingMode(false)
+
+		glog.Info("[import] Got proposal to get latest from leader P dir")
+		if proposal.Reqpdirstream.Addr == n.MyAddr {
+			glog.Info("[import] already have latest p dir skipping")
+			return nil
+		}
+
+		var err error
+		var closer *z.Closer
+		closer, err = n.startTask(opStreamPDir)
+		if err != nil {
+			return errors.Wrapf(err, "cannot start stream p dir task")
+		}
+		defer closer.Done()
+
+		pl, err := conn.GetPools().Get(proposal.Reqpdirstream.Addr)
+		if err != nil {
+			glog.Error("[import] unable to set connection pool")
+			return err
+		}
+
+		if pl == nil {
+			glog.Error("[import] connection is broken")
+			return fmt.Errorf("connection is broken")
+		}
+
+		con := pl.Get()
+		c := pb.NewWorkerClient(con)
+		r := &pb.ReqPDirStreamRequest{Addr: n.MyAddr}
+		stream, err := c.ReqPDirStream(ctx, r)
+		if err != nil {
+			glog.Error("[import] Error streaming P dir")
+			return err
+		}
+		return FlushData(stream)
+
 	case proposal.Snapshot != nil:
 		existing, err := n.Store.Snapshot()
 		if err != nil {