From 807f32573f31fd5d1cbc7efcaebaab9c6d934afd Mon Sep 17 00:00:00 2001
From: Eric Lin <linxiulei@gmail.com>
Date: Sat, 22 Jun 2024 09:15:46 +0000
Subject: [PATCH] cri: optimize ListPodSandboxStats with parallelism

Signed-off-by: Eric Lin <linxiulei@gmail.com>
---
 internal/cri/server/sandbox_stats_list.go | 38 ++++++++++++++++-------
 1 file changed, 26 insertions(+), 12 deletions(-)

diff --git a/internal/cri/server/sandbox_stats_list.go b/internal/cri/server/sandbox_stats_list.go
index 627026974ce0..e481ce35754e 100644
--- a/internal/cri/server/sandbox_stats_list.go
+++ b/internal/cri/server/sandbox_stats_list.go
@@ -20,6 +20,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"sync"
 
 	sandboxstore "github.com/containerd/containerd/v2/internal/cri/store/sandbox"
 	"github.com/containerd/errdefs"
@@ -34,20 +35,33 @@ func (c *criService) ListPodSandboxStats(
 	r *runtime.ListPodSandboxStatsRequest,
 ) (*runtime.ListPodSandboxStatsResponse, error) {
 	sandboxes := c.sandboxesForListPodSandboxStatsRequest(r)
+	stats, errs := make([]*runtime.PodSandboxStats, len(sandboxes)), make([]error, len(sandboxes))
+
+	var wg sync.WaitGroup
+	for i, sandbox := range sandboxes {
+		i := i
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			sandboxStats, err := c.podSandboxStats(ctx, sandbox)
+			switch {
+			case errdefs.IsUnavailable(err), errdefs.IsNotFound(err):
+				log.G(ctx).WithField("podsandboxid", sandbox.ID).WithError(err).Debug("failed to get pod sandbox stats, this is likely a transient error")
+			case errors.Is(err, ttrpc.ErrClosed):
+				log.G(ctx).WithField("podsandboxid", sandbox.ID).WithError(err).Debug("failed to get pod sandbox stats, connection closed")
+			case err != nil:
+				errs[i] = fmt.Errorf("failed to decode sandbox container metrics for sandbox %q: %w", sandbox.ID, err)
+			default:
+				stats[i] = sandboxStats
+			}
+		}()
+	}
+	wg.Wait()
 
-	var errs []error
 	podSandboxStats := new(runtime.ListPodSandboxStatsResponse)
-	for _, sandbox := range sandboxes {
-		sandboxStats, err := c.podSandboxStats(ctx, sandbox)
-		switch {
-		case errdefs.IsUnavailable(err), errdefs.IsNotFound(err):
-			log.G(ctx).WithField("podsandboxid", sandbox.ID).Debugf("failed to get pod sandbox stats, this is likely a transient error: %v", err)
-		case errors.Is(err, ttrpc.ErrClosed):
-			log.G(ctx).WithField("podsandboxid", sandbox.ID).Debugf("failed to get pod sandbox stats, connection closed: %v", err)
-		case err != nil:
-			errs = append(errs, fmt.Errorf("failed to decode sandbox container metrics for sandbox %q: %w", sandbox.ID, err))
-		default:
-			podSandboxStats.Stats = append(podSandboxStats.Stats, sandboxStats)
+	for _, stat := range stats {
+		if stat != nil {
+			podSandboxStats.Stats = append(podSandboxStats.Stats, stat)
 		}
 	}