perf: read workload statuses one prefix per app, entrypoint and node

store/common/workload.go

@@ -7,15 +7,17 @@ import (
 	"maps"
 	"path/filepath"
 	"slices"
-	"sync"
 
 	"github.com/cockroachdb/errors"
+	"golang.org/x/sync/errgroup"
 
 	"github.com/projecteru2/core/log"
 	"github.com/projecteru2/core/types"
 	"github.com/projecteru2/core/utils"
 )
 
+const statusReaders = 32
+
 func (s *Store) AddWorkload(ctx context.Context, workload *types.Workload, processing *types.Processing) error {
 	return s.doOpsWorkload(ctx, workload, processing, true)
 }
@@ -185,14 +187,16 @@ func (s *Store) filterWorkloads(ctx context.Context, data, labels map[string]str
 
 func (s *Store) bindWorkloadsAdditions(ctx context.Context, workloads []*types.Workload, withEngine bool) ([]*types.Workload, error) {
 	nodenames := map[string]struct{}{}
-	statusKeys := map[string]string{}
+	groups := map[string][]*types.Workload{}
 	logger := log.WithFunc("store.common.bindWorkloadsAdditions")
 	for _, workload := range workloads {
 		appname, entrypoint, _, err := utils.ParseWorkloadName(workload.Name)
 		if err != nil {
 			return nil, err
 		}
-		statusKeys[workload.ID] = filepath.Join(WorkloadStatusPrefix, appname, entrypoint, workload.Nodename, workload.ID)
+		// trailing slash keeps the prefix from matching a longer nodename
+		prefix := filepath.Join(WorkloadStatusPrefix, appname, entrypoint, workload.Nodename) + "/"
+		groups[prefix] = append(groups[prefix], workload)
 		nodenames[workload.Nodename] = struct{}{}
 	}
 	if withEngine {
@@ -213,24 +217,34 @@ func (s *Store) bindWorkloadsAdditions(ctx context.Context, workloads []*types.W
 		}
 	}
 
-	wg := &sync.WaitGroup{}
-	wg.Add(len(workloads))
-	for _, workload := range workloads {
-		_ = s.Pool.Invoke(func() {
-			defer wg.Done()
-			value, err := s.GetOne(ctx, statusKeys[workload.ID])
-			if err != nil {
-				return
+	bind := func(prefix string, group []*types.Workload) error {
+		data, err := s.GetPrefix(ctx, prefix, 0)
+		if err != nil {
+			return err
+		}
+		for _, workload := range group {
+			value, ok := data[prefix+workload.ID]
+			if !ok {
+				continue
 			}
 			status := &types.StatusMeta{}
-			if err := json.Unmarshal([]byte(value), &status); err != nil {
+			if err := json.Unmarshal([]byte(value), status); err != nil {
 				logger.Errorf(ctx, err, "unmarshal status of %s, raw: %s", workload.ID, value)
-				return
+				continue
 			}
 			workload.StatusMeta = status
-		})
+		}
+		return nil
+	}
+
+	reads := errgroup.Group{}
+	reads.SetLimit(statusReaders)
+	for prefix, group := range groups {
+		reads.Go(func() error { return bind(prefix, group) })
+	}
+	if err := reads.Wait(); err != nil {
+		return nil, err
 	}
-	wg.Wait()
 	return workloads, nil
 }
 

store/common/workload_test.go

@@ -1,6 +1,12 @@
 package common
 
 import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"slices"
+	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -33,3 +39,115 @@ func TestWorkloadStatusStreamClosesWhenWatchBreaks(t *testing.T) {
 		t.Fatal("workload status stream did not close after the watch broke")
 	}
 }
+
+func TestGetWorkloadsReadsOneStatusPrefixPerGroup(t *testing.T) {
+	kv := newStatusPrefixKV(t)
+	store := newStatusPrefixStore(t, kv)
+
+	workloads, err := store.getWorkloads(t.Context(), []string{"w1", "w2", "w3", "w4"}, false)
+	require.NoError(t, err)
+	require.Len(t, workloads, 4)
+
+	assert.Equal(t, []string{"/status/test/app/n1/", "/status/test/app/n2/", "/status/test/web/n1/"}, kv.readPrefixes())
+	require.NotNil(t, workloads[0].StatusMeta)
+	assert.True(t, workloads[0].StatusMeta.Running)
+	assert.Nil(t, workloads[1].StatusMeta, "a workload without a status key keeps a nil StatusMeta")
+	require.NotNil(t, workloads[2].StatusMeta)
+	assert.Equal(t, "w3", workloads[2].StatusMeta.ID)
+	assert.Nil(t, workloads[3].StatusMeta, "an unmarshalable status is skipped")
+}
+
+func TestGetWorkloadReadsOneStatusPrefix(t *testing.T) {
+	kv := newStatusPrefixKV(t)
+	store := newStatusPrefixStore(t, kv)
+
+	workload, err := store.getWorkload(t.Context(), "w1", false)
+	require.NoError(t, err)
+	assert.Equal(t, []string{"/status/test/app/n1/"}, kv.readPrefixes())
+	require.NotNil(t, workload.StatusMeta)
+	assert.Equal(t, "w1", workload.StatusMeta.ID)
+}
+
+func TestGetWorkloadsFailsWhenStatusReadFails(t *testing.T) {
+	kv := newStatusPrefixKV(t)
+	kv.err = types.ErrMockError
+	store := newStatusPrefixStore(t, kv)
+
+	_, err := store.getWorkloads(t.Context(), []string{"w1", "w3"}, false)
+	assert.ErrorIs(t, err, types.ErrMockError)
+}
+
+func newStatusPrefixStore(t *testing.T, kv KV) *Store {
+	t.Helper()
+	pool, err := utils.NewPool(1)
+	require.NoError(t, err)
+	t.Cleanup(pool.Release)
+	return New(kv, types.Config{}, pool)
+}
+
+func newStatusPrefixKV(t *testing.T) *statusPrefixKV {
+	t.Helper()
+	kv := &statusPrefixKV{workloads: map[string]string{}, statuses: map[string]string{}}
+	for _, workload := range []*types.Workload{
+		{ID: "w1", Name: "test_app_1", Nodename: "n1"},
+		{ID: "w2", Name: "test_app_2", Nodename: "n1"},
+		{ID: "w3", Name: "test_app_3", Nodename: "n2"},
+		{ID: "w4", Name: "test_web_1", Nodename: "n1"},
+	} {
+		data, err := json.Marshal(workload)
+		require.NoError(t, err)
+		kv.workloads[fmt.Sprintf(WorkloadInfoKey, workload.ID)] = string(data)
+	}
+	for key, ID := range map[string]string{
+		"/status/test/app/n1/w1": "w1",
+		"/status/test/app/n1/w9": "w9",
+		"/status/test/app/n2/w3": "w3",
+	} {
+		data, err := json.Marshal(&types.StatusMeta{ID: ID, Running: true})
+		require.NoError(t, err)
+		kv.statuses[key] = string(data)
+	}
+	kv.statuses["/status/test/web/n1/w4"] = "}"
+	return kv
+}
+
+type statusPrefixKV struct {
+	KV
+
+	workloads map[string]string
+	statuses  map[string]string
+	err       error
+
+	mu       sync.Mutex
+	prefixes []string
+}
+
+func (k *statusPrefixKV) GetMulti(_ context.Context, keys []string) (map[string]string, error) {
+	data := make(map[string]string, len(keys))
+	for _, key := range keys {
+		data[key] = k.workloads[key]
+	}
+	return data, nil
+}
+
+func (k *statusPrefixKV) GetPrefix(_ context.Context, prefix string, _ int64) (map[string]string, error) {
+	k.mu.Lock()
+	k.prefixes = append(k.prefixes, prefix)
+	k.mu.Unlock()
+	if k.err != nil {
+		return nil, k.err
+	}
+	data := map[string]string{}
+	for key, value := range k.statuses {
+		if strings.HasPrefix(key, prefix) {
+			data[key] = value
+		}
+	}
+	return data, nil
+}
+
+func (k *statusPrefixKV) readPrefixes() []string {
+	k.mu.Lock()
+	defer k.mu.Unlock()
+	return slices.Sorted(slices.Values(k.prefixes))
+}