add stream list concurrency config

vdaas · Aug 18, 2023 · 04b5ad1 · 04b5ad1
1 parent b67c272
commit 04b5ad1
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 2 deletions.
diff --git a/cmd/index/job/correction/sample.yaml b/cmd/index/job/correction/sample.yaml
@@ -78,6 +78,7 @@ corrector:
   agent_dns: vald-agent-ngt.default.svc.cluster.local
   agent_namespace: "default"
   node_name: ""
+  stream_list_concurrency: 100
   discoverer:
     duration: 500ms
     client:

diff --git a/internal/config/corrector.go b/internal/config/corrector.go
@@ -36,6 +36,10 @@ type Corrector struct {
 	// NodeName represents node name
 	NodeName string `json:"node_name" yaml:"node_name"`
 
+	// StreamConcurrency represent stream concurrency for StreamListObject rpc client
+	// this directly affects the memory usage of this job
+	StreamListConcurrency int `json:"stream_list_concurrency" yaml:"stream_list_concurrency"`
+
 	// Discoverer represent agent discoverer service configuration
 	Discoverer *DiscovererClient `json:"discoverer" yaml:"discoverer"`
 }
@@ -52,3 +56,11 @@ func (c *Corrector) Bind() *Corrector {
 	}
 	return c
 }
+
+// GetStreamListConcurrency returns the StreamListConcurrency field value if set, -1 otherwise, which means no limit.
+func (c *Corrector) GetStreamListConcurrency() int {
+	if c != nil {
+		return c.StreamListConcurrency
+	}
+	return -1
+}
diff --git a/pkg/index/job/correction/service/corrector.go b/pkg/index/job/correction/service/corrector.go
@@ -117,7 +117,10 @@ func (c *correct) correct(ctx context.Context, addrs []string) (err error) {
 			}
 
 			seg, ctx := stdeg.WithContext(ctx)
-			seg.SetLimit(c.cfg.Server.GetGRPCStreamConcurrency()) // FIXME: server settingsをそのまま流用で良いのか？
+			concurrency := c.cfg.Corrector.GetStreamListConcurrency()
+			seg.SetLimit(concurrency) // FIXME: server settingsをそのまま流用で良いのか？
+
+			log.Infof("starting correction for agent %s, concurrency: %d", addr, concurrency)
 
 			finalize := func() error {
 				err = seg.Wait()
@@ -205,7 +208,7 @@ func (c *correct) correct2(ctx context.Context, addrs []string) (err error) {
 	if err := c.discoverer.GetClient().OrderedRange(ctx, addrs,
 		func(ctx context.Context, addr string, conn *grpc.ClientConn, copts ...grpc.CallOption) error {
 			eg, ctx := errgroup.New(ctx)
-			eg.Limitation(c.cfg.Server.GetGRPCStreamConcurrency())
+			eg.Limitation(c.cfg.Corrector.GetStreamListConcurrency())
 
 			vc := vald.NewValdClient(conn)
 			stream, err := vc.StreamListObject(ctx, &payload.Object_List_Request{})