From 884c9f950717778734571dae28812f94ffeacbb8 Mon Sep 17 00:00:00 2001 From: CMGS Date: Thu, 10 Sep 2026 01:05:51 +0900 Subject: [PATCH] perf: read workload statuses one prefix per app, entrypoint and node bindWorkloadsAdditions fetched every workload's status with its own GetOne; a listing of N workloads cost N status round trips on top of the meta and node reads. Workloads sharing an appname, entrypoint and node share a status key prefix, so one GetPrefix per group replaces them: N*M reads become M for M nodes. Groups read concurrently under an errgroup capped at 32. A status read failure now fails the call instead of leaving the status nil silently; a missing or unmarshalable status still leaves StatusMeta nil. --- store/common/workload.go | 44 ++++++++----- store/common/workload_test.go | 118 ++++++++++++++++++++++++++++++++++ 2 files changed, 147 insertions(+), 15 deletions(-) diff --git a/store/common/workload.go b/store/common/workload.go index e24948da..187c9187 100644 --- a/store/common/workload.go +++ b/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 } diff --git a/store/common/workload_test.go b/store/common/workload_test.go index 9b585479..41725dc8 100644 --- a/store/common/workload_test.go +++ b/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)) +}