Add support for HA and multishard functionality in import APIs #9406

shivaji-kharse · 2025-05-13T12:37:12Z

Description

Please explain the changes you made here.

Checklist

Code compiles correctly and linting passes locally
For all code changes, an entry added to the CHANGELOG.md file describing and linking to
this PR
Tests added for new functionality, or regression tests for bug fixes added as applicable
For public APIs, new features, etc., PR on
docs repo staged and linked here

Instructions

The PR title should follow the Conventional Commits
syntax, leading with fix:, feat:, chore:, ci:, etc.
The description should briefly explain what the PR is about. In the case of a bugfix, describe or
link to the bug.
In the checklist section, check the boxes in that are applicable, using [x] syntax.
- If not applicable, remove the entire line. Only leave the box unchecked if you intend to come
  back and check the box later.
Delete the Instructions line and everything below it, to indicate you have read and are
following these instructions. 🙂

Thank you for your contribution to Dgraph!

Copilot

Pull Request Overview

This PR adds support for high availability and multishard functionality in the import APIs by introducing a new streaming mechanism for propagating partition directory (P dir) updates across nodes. Key changes include:

Adding new stream processing functions and a helper for managing ongoing P dir tasks.
Introducing a new RPC (ReqPDirStream) and related message and client/server implementations.
Updating tests and client code to leverage the new streaming workflow.

Reviewed Changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
worker/import.go	Adds streamProcessorSender support, a new stream task check, and updates the acknowledgment flow for stream processing.
worker/draft.go	Introduces a new opStreamPDir constant and string conversion for task identification.
protos/pb/pb_grpc.pb.go	Implements the new ReqPDirStream RPC handler methods and client stream handling.
protos/pb.proto	Defines the new ReqPDirStreamRequest message and RPC in the Worker service.
edgraph/server.go	Removes legacy drain mode invocation in favor of the new stream-based approach.
dgraph/cmd/dgraphimport/import_test.go	Updates test parameters and introduces a longer fixed sleep period post-import.
dgraph/cmd/dgraphimport/import_client.go	Refactors stream data logic to use the new RunBadgerStream method from the worker package.

Comments suppressed due to low confidence (1)

worker/import.go:319

The removal of sending the acknowledgment using stream.SendAndClose in this function appears to be intentional given the updated flow in InStream. Please verify that the ACK signal is being sent exactly once during stream processing to avoid protocol inconsistencies.

return nil

Copilot · 2025-05-13T12:37:50Z

dgraph/cmd/dgraphimport/import_test.go

+	time.Sleep(time.Minute * 3)
 	verifyImportResults(t, gc)


[nitpick] Using a fixed sleep of 3 minutes may unnecessarily prolong test execution. Consider replacing the fixed delay with a synchronization mechanism or a condition-based wait to reduce test runtime.

Suggested change

time.Sleep(time.Minute * 3)

verifyImportResults(t, gc)

ctx, cancel := context.WithTimeout(context.Background(), time.Minute*3)

defer cancel()

for {

if verifyImportResults(t, gc) {

break

}

select {

case <-ctx.Done():

t.Fatal("Timeout waiting for import results to be ready")

default:

time.Sleep(time.Second * 5) // Poll every 5 seconds

}

}

trunk-io · 2025-05-13T13:34:02Z

Failed Test	Failure Summary	Logs
`TestVectorIndexWithoutSchema`	Error occurred while computing euclidean distance on vectors of different lengths.	Logs ↗︎

_{View Full Report ↗︎ ⋅ Docs}

trunk-io · 2025-05-26T13:01:54Z

Running Code Quality on PRs by uploading data to Trunk will soon be removed. You can still run checks on your PRs using trunk-action - see the migration guide for more information.

dgraph/cmd/dgraphimport/import_client.go

mangalaman93 · 2025-05-29T10:56:01Z

dgraph/cmd/dgraphimport/import_client.go

 	}
-	// Wait for acknowledgment from the server
 	if _, err := out.CloseAndRecv(); err != nil {


do we need to close the steam even if there is an error before this line

dgraph/cmd/dgraphimport/import_test.go

mangalaman93 · 2025-05-29T11:38:24Z

dgraph/cmd/dgraphimport/import_test.go

+	require.NoError(t, targetCluster.HealthCheck(false))
+
+	if waitForSnapshot {
+		for grp, alphas := range alphaGroups {


why are we waiting for a snapshot?

So when a stopped Alpha rejoins the group, it receives a Raft proposal and then requests the latest snapshot from the leader. We should wait for that process to complete — that's why we are waiting for the snapshot.

mangalaman93 · 2025-05-29T11:40:17Z

x/health.go

 	"github.com/pkg/errors"
 )

 var (
 	// the drainingMode variable should be accessed through the atomic.Store and atomic.Load
 	// functions. The value 0 means the draining-mode is disabled, and the value 1 means the
 	// mode is enabled
-	drainingMode uint32
+	drainingMode              uint32
+	extSnapshotStreamingState uint32


do we need this?

Yes, so we don't want to rely on drain mode only, because it can be changed by other operations as well.

worker/draft.go

Copilot AI review requested due to automatic review settings May 13, 2025 12:37

shivaji-kharse requested a review from a team as a code owner May 13, 2025 12:37

github-actions bot added area/testing Testing related issues area/core internal mechanisms go Pull requests that update Go code labels May 13, 2025

Copilot AI reviewed May 13, 2025

View reviewed changes

shivaji-kharse force-pushed the shiva/stream-ha branch from 2572f94 to dc05732 Compare May 26, 2025 13:01

shivaji-kharse force-pushed the shiva/stream-ha branch 3 times, most recently from f2a9fed to 506301a Compare May 28, 2025 06:52

Add support for HA and multishard functionality in import APIs

eddd33c

shivaji-kharse force-pushed the shiva/stream-ha branch from ada9899 to eddd33c Compare May 29, 2025 09:29

shivaji-kharse added 2 commits May 29, 2025 15:13

fix trunk

b1b96ea

fix tests

1f2eaad

mangalaman93 requested changes May 29, 2025

View reviewed changes

resolve review comments

687d521

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add support for HA and multishard functionality in import APIs #9406

Add support for HA and multishard functionality in import APIs #9406

shivaji-kharse commented May 13, 2025

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI May 13, 2025

Uh oh!

trunk-io bot commented May 13, 2025 •

edited

Loading

Uh oh!

trunk-io bot commented May 26, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mangalaman93 May 29, 2025

Uh oh!

Uh oh!

Uh oh!

mangalaman93 May 29, 2025

Uh oh!

shivaji-kharse May 29, 2025

Uh oh!

mangalaman93 May 29, 2025

Uh oh!

shivaji-kharse May 29, 2025

Uh oh!

Uh oh!

Uh oh!

-	time.Sleep(time.Minute * 3)
-	verifyImportResults(t, gc)
+	ctx, cancel := context.WithTimeout(context.Background(), time.Minute*3)
+	defer cancel()
+	for {
+		if verifyImportResults(t, gc) {
+			break
+		}
+		select {
+		case <-ctx.Done():
+			t.Fatal("Timeout waiting for import results to be ready")
+		default:
+			time.Sleep(time.Second * 5) // Poll every 5 seconds
+		}
+	}

Add support for HA and multishard functionality in import APIs #9406

Are you sure you want to change the base?

Add support for HA and multishard functionality in import APIs #9406

Conversation

shivaji-kharse commented May 13, 2025

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Copilot AI May 13, 2025

Choose a reason for hiding this comment

Uh oh!

trunk-io bot commented May 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

trunk-io bot commented May 26, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mangalaman93 May 29, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

mangalaman93 May 29, 2025

Choose a reason for hiding this comment

Uh oh!

shivaji-kharse May 29, 2025

Choose a reason for hiding this comment

Uh oh!

mangalaman93 May 29, 2025

Choose a reason for hiding this comment

Uh oh!

shivaji-kharse May 29, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

trunk-io bot commented May 13, 2025 •

edited

Loading