Refactor

refactor

pkg/index/job/correction/service/corrector.go

@@ -34,6 +34,7 @@ import (
 	"github.com/vdaas/vald/internal/net/grpc"
 	"github.com/vdaas/vald/internal/net/grpc/codes"
 	"github.com/vdaas/vald/internal/net/grpc/status"
+	"github.com/vdaas/vald/internal/safety"
 	"github.com/vdaas/vald/internal/sync"
 	"github.com/vdaas/vald/internal/sync/errgroup"
 	"github.com/vdaas/vald/pkg/index/job/correction/config"
@@ -106,9 +107,8 @@ func (c *correct) Start(ctx context.Context) (<-chan error, error) {
 		return nil, err
 	}
 
-	// For debugging
 	c.indexInfos.Range(func(addr string, info *payload.Info_Index_Count) bool {
-		log.Debugf("index info: addr(%s), stored(%d), uncommitted(%d)", addr, info.GetStored(), info.GetUncommitted())
+		log.Infof("index info: addr(%s), stored(%d), uncommitted(%d)", addr, info.GetStored(), info.GetUncommitted())
 		return true
 	})
 
@@ -133,30 +133,22 @@ func (c *correct) correct(ctx context.Context) (err error) {
 	// This is used to know which agents possibly have the same index as the target replica.
 	// We can say this because, thanks to caching, there is no way that the target replica is
 	// in the agent that has already been corrected.
-	leftAgentAddrs := make([]string, len(c.agentAddrs))
-	n := copy(leftAgentAddrs, c.agentAddrs)
-	if n != len(c.agentAddrs) {
-		return fmt.Errorf("failed to copy agentAddrs")
-	}
 
 	// Vector with time after this should not be processed
-	correctionStartTime, err := getCorrectionStartTime(ctx)
+	correctionStartTime, err := correctionStartTime(ctx)
 	if err != nil {
 		log.Errorf("cannot determine correction start time: %w", err)
 		return err
 	}
 
+	curTargetAgent := 0
 	if err := c.discoverer.GetClient().OrderedRange(ctx, c.agentAddrs,
 		func(ctx context.Context, addr string, conn *grpc.ClientConn, copts ...grpc.CallOption) error {
 			// current address is the leftAgentAddrs[0] because this is OrderedRange and
 			// leftAgentAddrs is copied from c.agentAddrs
-			leftAgentAddrs = leftAgentAddrs[1:]
-
-			vc := vald.NewValdClient(conn)
-			stream, err := vc.StreamListObject(ctx, &payload.Object_List_Request{})
-			if err != nil {
-				return err
-			}
+			defer func() {
+				curTargetAgent++
+			}()
 
 			// context and errgroup for stream.Recv and correction
 			sctx, scancel := context.WithCancel(ctx)
@@ -173,6 +165,12 @@ func (c *correct) correct(ctx context.Context) (err error) {
 			var mu sync.Mutex
 			log.Infof("starting correction for agent %s, stream concurrency: %d, bbolt concurrency: %d", addr, sconcurrency, bconcurrency)
 
+			vc := vald.NewValdClient(conn)
+			stream, err := vc.StreamListObject(ctx, &payload.Object_List_Request{})
+			if err != nil {
+				return err
+			}
+
 			// The number of items to be received in advance is not known in advance.
 			// This is because there is a possibility of new items being inserted during processing.
 			for {
@@ -181,9 +179,26 @@ func (c *correct) correct(ctx context.Context) (err error) {
 					if !errors.Is(sctx.Err(), context.Canceled) {
 						log.Errorf("context done unexpectedly: %v", sctx.Err())
 					}
-					goto Finalize
+
+					// Finalize
+					err = seg.Wait()
+					if err != nil {
+						log.Errorf("err group returned error: %v", err)
+					}
+
+					berr := bolteg.Wait()
+					if berr != nil {
+						log.Errorf("bbolt err group returned error: %v", err)
+						err = errors.Join(err, berr)
+					} else {
+						log.Info("bbolt all batch finished")
+					}
+
+					log.Infof("correction finished for agent %s", addr)
+					return err
+
 				default:
-					seg.Go(func() error {
+					seg.Go(safety.RecoverFunc(func() error {
 						mu.Lock()
 						// As long as we don't stream.Recv() from the stream, we do not consume the memory of the message.
 						// So by limiting the number of this errgroup.Go instances, we can limit the memory usage
@@ -235,7 +250,7 @@ func (c *correct) correct(ctx context.Context) (err error) {
 								addr: addr,
 								vec:  vec,
 							},
-							leftAgentAddrs,
+							curTargetAgent,
 						); err != nil {
 							log.Errorf("failed to check consistency: %v", err)
 							return nil // continue other processes
@@ -245,25 +260,9 @@ func (c *correct) correct(ctx context.Context) (err error) {
 						c.checkedID.AsyncSet(bolteg, []byte(id), nil)
 
 						return nil
-					})
+					}))
 				}
 			}
-
-		Finalize:
-			err = seg.Wait()
-			if err != nil {
-				log.Errorf("err group returned error: %v", err)
-			}
-
-			berr := bolteg.Wait()
-			if berr != nil {
-				log.Errorf("bolt err group returned error: %v", err)
-				err = errors.Join(err, berr)
-			}
-			log.Info("bbolt all batch finished")
-
-			log.Infof("correction finished for agent %s", addr)
-			return err
 		},
 	); err != nil {
 		log.Errorf("failed to range over agents(%v): %v", c.agentAddrs, err)
@@ -279,39 +278,23 @@ type vectorReplica struct {
 }
 
 // Validate len(addrs) >= 2 before calling this function
-func (c *correct) checkConsistency(ctx context.Context, targetReplica *vectorReplica, leftAgentAddrs []string) error {
-	// availableAddrs is the agents' addr that doesn't have the target replica thus is available to insert the replica
-	// to fix the index replica number if required.
-	availableAddrs := make([]string, 0, len(c.agentAddrs)-1)
-	for _, addr := range c.agentAddrs {
-		if addr != targetReplica.addr {
-			availableAddrs = append(availableAddrs, addr)
-		}
-	}
+// idxだけ渡せば良い？c.addrsに全ての情報があるので？
+func (c *correct) checkConsistency(ctx context.Context, targetReplica *vectorReplica, targetAgentIdx int) error {
+	// leftAgentAddrs is the agents' addr that hasn't been corrected yet.
+	leftAgentAddrs := c.agentAddrs[targetAgentIdx+1:]
 
 	// Vector with time after this should not be processed
-	correctionStartTime, err := getCorrectionStartTime(ctx)
+	correctionStartTime, err := correctionStartTime(ctx)
 	if err != nil {
 		log.Errorf("cannot determine correction start time: %w", err)
 		return err
 	}
 
-	foundReplicas := make([]*vectorReplica, 0, len(availableAddrs))
+	foundReplicas := make([]*vectorReplica, 0, len(c.agentAddrs))
 	var mu sync.Mutex
 	if err := c.discoverer.GetClient().OrderedRangeConcurrent(ctx, leftAgentAddrs, len(leftAgentAddrs),
 		func(ctx context.Context, addr string, conn *grpc.ClientConn, copts ...grpc.CallOption) error {
-			// To avoid GetObject to myself. To maintain backward compatibility for withoug cache operation
-			if addr == targetReplica.addr {
-				return nil
-			}
-
-			select {
-			case <-ctx.Done():
-				return ctx.Err()
-			default:
-			}
-			vc := vald.NewValdClient(conn)
-			vec, err := vc.GetObject(ctx, &payload.Object_VectorRequest{
+			vec, err := vald.NewValdClient(conn).GetObject(ctx, &payload.Object_VectorRequest{
 				Id: &payload.Object_ID{
 					Id: targetReplica.vec.GetId(),
 				},
@@ -344,12 +327,6 @@ func (c *correct) checkConsistency(ctx context.Context, targetReplica *vectorRep
 				addr: addr,
 				vec:  vec,
 			})
-
-			// Remove this addr from availableAddrs because this addr has the target replica
-			// and not available to insert the replica to fix the index replica number
-			slices.DeleteFunc(availableAddrs, func(availableAddr string) bool {
-				return availableAddr == addr
-			})
 			mu.Unlock()
 
 			return nil
@@ -364,7 +341,7 @@ func (c *correct) checkConsistency(ctx context.Context, targetReplica *vectorRep
 	}
 
 	// check replica number
-	if err := c.correctReplica(ctx, targetReplica, foundReplicas, availableAddrs); err != nil {
+	if err := c.correctReplica(ctx, targetReplica, foundReplicas); err != nil {
 		return fmt.Errorf("failed to fix index replica: %w", err)
 	}
 
@@ -413,7 +390,6 @@ func (c *correct) correctReplica(
 	ctx context.Context,
 	targetReplica *vectorReplica,
 	foundReplicas []*vectorReplica,
-	availableAddrs []string,
 ) error {
 	// diff < 0 means there is less replica than the correct number
 	existReplica := len(foundReplicas) + 1
@@ -423,6 +399,19 @@ func (c *correct) correctReplica(
 		return nil
 	}
 
+	availableAddrs := make([]string, 0, len(c.agentAddrs))
+	for _, addr := range c.agentAddrs {
+		if addr == targetReplica.addr {
+			continue
+		}
+		if slices.ContainsFunc(foundReplicas, func(replica *vectorReplica) bool {
+			return replica.addr == addr
+		}) {
+			continue
+		}
+		availableAddrs = append(availableAddrs, addr)
+	}
+
 	// when there are less replicas than the correct number, add the extra replicas
 	if diff < 0 {
 		log.Infof("replica shortage of vector %s. inserting to other agents...", targetReplica.vec.GetId())
@@ -596,7 +585,7 @@ func embedTime(ctx context.Context) context.Context {
 	return context.WithValue(ctx, correctionStartTimeKey, time.Now())
 }
 
-func getCorrectionStartTime(ctx context.Context) (time.Time, error) {
+func correctionStartTime(ctx context.Context) (time.Time, error) {
 	v := ctx.Value(correctionStartTimeKey)
 	if t, ok := v.(time.Time); ok {
 		return t, nil

pkg/index/job/correction/service/corrector_test.go

@@ -437,9 +437,16 @@ func Test_correct_correctReplica(t *testing.T) {
 				},
 			},
 		}
+
+		// agentAddrs = availableAddrs + target.addr + found.addr
+		c.agentAddrs = append(test.args.availableAddrs, test.args.target.addr)
+		for _, found := range test.args.found {
+			c.agentAddrs = append(c.agentAddrs, found.addr)
+		}
+
 		t.Run(test.name, func(tt *testing.T) {
 			tt.Parallel()
-			err := c.correctReplica(context.Background(), test.args.target, test.args.found, test.args.availableAddrs)
+			err := c.correctReplica(context.Background(), test.args.target, test.args.found)
 			if test.want.err != nil {
 				require.ErrorIs(t, test.want.err, err)
 			}