@@ -12,14 +12,13 @@ import (
1212
1313 apiv25 "github.com/dgraph-io/dgo/v250/protos/api.v25"
1414 "github.com/dgraph-io/ristretto/v2/z"
15- "github.com/dustin/go-humanize"
16- "github.com/golang/glog"
1715 "github.com/hypermodeinc/dgraph/v25/conn"
1816 "github.com/hypermodeinc/dgraph/v25/posting"
19-
2017 "github.com/hypermodeinc/dgraph/v25/protos/pb"
21-
2218 "github.com/hypermodeinc/dgraph/v25/schema"
19+
20+ "github.com/dustin/go-humanize"
21+ "github.com/golang/glog"
2322 "github.com/pkg/errors"
2423)
2524
@@ -43,8 +42,7 @@ func ProposeDrain(ctx context.Context, drainMode *pb.Drainmode) error {
4342 }
4443 con := pl .Get ()
4544 c := pb .NewWorkerClient (con )
46- _ , err := c .ApplyDrainmode (ctx , drainMode )
47- if err != nil {
45+ if _ , err := c .ApplyDrainmode (ctx , drainMode ); err != nil {
4846 return err
4947 }
5048 }
@@ -54,18 +52,18 @@ func ProposeDrain(ctx context.Context, drainMode *pb.Drainmode) error {
5452
5553// DoStreamPDir handles streaming of snapshots to a target group. It first checks the group
5654// associated with the incoming stream and, if it's the same as the current node's group, it
57- // flushes the data using FlushKvs1 . If the group is different, it establishes a connection
55+ // flushes the data using FlushKvs . If the group is different, it establishes a connection
5856// with the leader of that group and streams data to it. The function returns an error if
5957// there are any issues in the process, such as a broken connection or failure to establish
6058// a stream with the leader.
6159func DoStreamPDir (stream apiv25.Dgraph_StreamSnapshotServer ) error {
62- groupId , err := checkGroup (stream )
60+ groupId , err := getGroup (stream )
6361 if err != nil {
6462 return err
6563 }
6664
6765 if groupId == groups ().Node .gid {
68- return FlushKvs (stream )
66+ return flushKvs (stream )
6967 }
7068
7169 pl := groups ().Leader (groupId )
@@ -82,46 +80,67 @@ func DoStreamPDir(stream apiv25.Dgraph_StreamSnapshotServer) error {
8280
8381 return streamToAnotherGroup (stream , out )
8482}
85-
8683func streamToAnotherGroup (in apiv25.Dgraph_StreamSnapshotServer , out pb.Worker_StreamPDirClient ) error {
8784 chan1 := make (chan * apiv25.StreamSnapshotRequest , 10 )
85+ errCh := make (chan error , 1 )
86+ ctx := in .Context ()
8887
8988 go func () {
9089 defer close (chan1 )
9190 for {
92- msg , err := in .Recv ()
93- if err != nil {
94- if err != io .EOF {
95- glog .Errorf ("Error receiving from in stream: %v" , err )
96- }
91+ select {
92+ case <- ctx .Done ():
93+ glog .Infof ("Context cancelled, stopping receive goroutine." )
94+ errCh <- fmt .Errorf ("context deadline exceeded" )
9795 return
96+ default :
97+ msg , err := in .Recv ()
98+ if err != nil {
99+ if err != io .EOF {
100+ glog .Errorf ("Error receiving from in stream: %v" , err )
101+ errCh <- err
102+ }
103+ return
104+ }
105+ chan1 <- msg
98106 }
99- chan1 <- msg
100107 }
101108 }()
102109
103110 size := 0
104111
105- for msg := range chan1 {
106- data := & pb.KVS {Data : msg .Pairs .Data }
112+ Loop:
113+ for {
114+ select {
115+ case err := <- errCh :
116+ return err
117+
118+ case msg , ok := <- chan1 :
119+ if ! ok {
120+ // Channel closed, exit loop
121+ break Loop
122+ }
123+
124+ data := & pb.KVS {Data : msg .Pairs .Data }
125+
126+ if msg .Pairs .Done {
127+ if err := out .Send (& pb.KVS {Done : true }); err != nil {
128+ glog .Errorf ("Error sending 'done' to out stream: %v" , err )
129+ return err
130+ }
131+ glog .Infoln ("All key-values have been transferred." )
132+ break Loop
133+ }
107134
108- if msg .Pairs .Done {
109- if err := out .Send (& pb.KVS {Done : true }); err != nil {
110- glog .Errorf ("Error sending 'done' to out stream: %v" , err )
135+ if err := out .Send (data ); err != nil {
136+ glog .Errorf ("Error sending to outstream: %v" , err )
111137 return err
112138 }
113- glog .Infoln ("All key-values have been transferred." )
114- break
115- }
116139
117- if err := out . Send ( data ); err != nil {
118- glog .Errorf ( "Error sending to outstream : %v" , err )
119- return err
140+ size += len ( msg . Pairs . Data )
141+ glog .Infof ( "Sent batch of size : %s. Total so far: %s \n " ,
142+ humanize . IBytes ( uint64 ( len ( msg . Pairs . Data ))), humanize . IBytes ( uint64 ( size )))
120143 }
121-
122- size += len (msg .Pairs .Data )
123- glog .Infof ("Sent batch of size: %s. Total so far: %s\n " ,
124- humanize .IBytes (uint64 (len (msg .Pairs .Data ))), humanize .IBytes (uint64 (size )))
125144 }
126145
127146 // Close the incoming stream properly
@@ -136,14 +155,13 @@ func streamToAnotherGroup(in apiv25.Dgraph_StreamSnapshotServer, out pb.Worker_S
136155 }
137156
138157 glog .Infof ("Received ACK with message: %v\n " , ack .Done )
139-
140158 return nil
141159}
142160
143- // checkGroup receives the initial message from the stream and extracts the group ID.
161+ // getGroup receives the initial message from the stream and extracts the group ID.
144162// It returns the group ID if successful, otherwise an error if there is an issue
145163// receiving the message.
146- func checkGroup (stream apiv25.Dgraph_StreamSnapshotServer ) (uint32 , error ) {
164+ func getGroup (stream apiv25.Dgraph_StreamSnapshotServer ) (uint32 , error ) {
147165 req , err := stream .Recv ()
148166 if err != nil {
149167 return 0 , fmt .Errorf ("failed to receive initial stream message: %v" , err )
@@ -152,9 +170,9 @@ func checkGroup(stream apiv25.Dgraph_StreamSnapshotServer) (uint32, error) {
152170 return req .GroupId , nil
153171}
154172
155- // FlushKvs receives the stream of data from the client and writes it to BadgerDB.
173+ // flushKvs receives the stream of data from the client and writes it to BadgerDB.
156174// It also sends a streams the data to other nodes of the same group and reloads the schema from the DB.
157- func FlushKvs (stream apiv25.Dgraph_StreamSnapshotServer ) error {
175+ func flushKvs (stream apiv25.Dgraph_StreamSnapshotServer ) error {
158176 var writer badgerWriter
159177 sw := pstore .NewStreamWriter ()
160178 defer sw .Cancel ()
0 commit comments