From 4bd7796232499d5f4de5338d290648496d5225d3 Mon Sep 17 00:00:00 2001
From: CMGS <ilskdw@gmail.com>
Date: Tue, 29 Sep 2026 14:42:20 +0800
Subject: [PATCH] rpc: answer GetWorkload for a missing workload with
 codes.NotFound

Clients had to string-match the error text to tell a missing workload from a failure (agent's WorkloadExists). GetWorkload now returns the standard NotFound code when the error chain holds ErrWorkloadNotExists; other failures keep the per-RPC code. The ErrInvaildCount mark on the etcd not-found error added in #756 is dropped: nothing reads it (resource-extend's nodestore uses the raw etcd client and its own ErrNodeNotExists).
---
 rpc/rpc.go                     |  3 +++
 rpc/rpc_test.go                | 17 +++++++++++++++++
 store/etcdv3/meta/etcd.go      |  4 ++--
 store/etcdv3/meta/etcd_test.go |  3 +--
 4 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/rpc/rpc.go b/rpc/rpc.go
index 9b7c5366..ab4da692 100644
--- a/rpc/rpc.go
+++ b/rpc/rpc.go
@@ -378,6 +378,9 @@ func (v *Vibranium) GetWorkload(ctx context.Context, ID *pb.WorkloadID) (*pb.Wor
 	task := v.newTask(ctx, "GetWorkload", false)
 	defer task.done()
 	workload, err := v.cluster.GetWorkload(task.context, ID.Id)
+	if errors.Is(err, types.ErrWorkloadNotExists) {
+		return nil, grpcstatus.Error(codes.NotFound, err.Error())
+	}
 	if err != nil {
 		return nil, grpcstatus.Error(GetWorkload, err.Error())
 	}
diff --git a/rpc/rpc_test.go b/rpc/rpc_test.go
index 223e1ad3..415036ce 100644
--- a/rpc/rpc_test.go
+++ b/rpc/rpc_test.go
@@ -2,11 +2,13 @@ package rpc
 
 import (
 	"context"
+	"errors"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
 	grpcstatus "google.golang.org/grpc/status"
 
 	grpcmocks "github.com/projecteru2/core/3rdmocks"
@@ -198,6 +200,21 @@ func TestRemoveWorkloadReportsItsOwnStatusCode(t *testing.T) {
 	assert.Equal(t, RemoveWorkload, grpcstatus.Code(err))
 }
 
+func TestGetWorkloadReportsAMissingWorkloadAsNotFound(t *testing.T) {
+	v := newVibranium()
+
+	cluster := v.cluster.(*clustermock.Cluster)
+	cluster.On("GetWorkload", mock.Anything, "gone").
+		Return(nil, errors.Join(types.ErrWorkloadNotExists, types.ErrKeyNotFound)).Once()
+	cluster.On("GetWorkload", mock.Anything, "broken").
+		Return(nil, types.ErrMockError).Once()
+
+	_, err := v.GetWorkload(context.Background(), &pb.WorkloadID{Id: "gone"})
+	assert.Equal(t, codes.NotFound, grpcstatus.Code(err))
+	_, err = v.GetWorkload(context.Background(), &pb.WorkloadID{Id: "broken"})
+	assert.Equal(t, GetWorkload, grpcstatus.Code(err))
+}
+
 func TestSendLargeFileReportsItsOwnStatusCode(t *testing.T) {
 	v := newVibranium()
 
diff --git a/store/etcdv3/meta/etcd.go b/store/etcdv3/meta/etcd.go
index b2a6a1a9..931c8637 100644
--- a/store/etcdv3/meta/etcd.go
+++ b/store/etcdv3/meta/etcd.go
@@ -115,7 +115,7 @@ func (e *ETCD) GetOne(ctx context.Context, key string, opts ...clientv3.OpOption
 		return nil, err
 	}
 	if resp.Count == 0 {
-		return nil, errors.Wrapf(errors.Mark(types.ErrKeyNotFound, types.ErrInvaildCount), "key: %s", key)
+		return nil, errors.Wrapf(types.ErrKeyNotFound, "key: %s", key)
 	}
 	if resp.Count != 1 {
 		return nil, errors.Wrapf(types.ErrInvaildCount, "key: %s", key)
@@ -135,7 +135,7 @@ func (e *ETCD) GetMulti(ctx context.Context, keys []string) ([]*mvccpb.KeyValue,
 	for idx, responseOp := range txnResponse.Responses {
 		resp := responseOp.GetResponseRange()
 		if resp.Count == 0 {
-			return nil, errors.Wrapf(errors.Mark(types.ErrKeyNotFound, types.ErrInvaildCount), "key: %s", keys[idx])
+			return nil, errors.Wrapf(types.ErrKeyNotFound, "key: %s", keys[idx])
 		}
 		if resp.Count != 1 {
 			return nil, errors.Wrapf(types.ErrInvaildCount, "key: %s", keys[idx])
diff --git a/store/etcdv3/meta/etcd_test.go b/store/etcdv3/meta/etcd_test.go
index a615ff70..01608baa 100644
--- a/store/etcdv3/meta/etcd_test.go
+++ b/store/etcdv3/meta/etcd_test.go
@@ -7,7 +7,6 @@ import (
 	"sync"
 	"testing"
 
-	"github.com/cockroachdb/errors"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 	"go.etcd.io/etcd/api/v3/etcdserverpb"
@@ -41,7 +40,7 @@ func TestGetOneMissingKeyIsNotFound(t *testing.T) {
 	e := NewEmbeddedETCD(t)
 	kv, err := e.GetOne(t.Context(), "/absent")
 	require.ErrorIs(t, err, types.ErrKeyNotFound)
-	require.True(t, errors.Is(err, types.ErrInvaildCount))
+	require.NotErrorIs(t, err, types.ErrInvaildCount)
 	require.EqualError(t, err, "key: /absent: key not found")
 	require.Nil(t, kv)
 }