Use principal to classify system workers with task queue prefix fallback

When a principal is available from the gRPC context, use its type to
determine if a worker is a system worker (type="temporal"). When
principal is not available (propagation disabled), fall back to the
existing task queue prefix check.

Also fixes isSystemWorkerWorker typo from previous rename.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: rkannan82

service/matching/handler.go

@@ -596,12 +596,13 @@ func (h *Handler) ListNexusEndpoints(ctx context.Context, request *matchingservi
 
 // RecordWorkerHeartbeat receive heartbeat request from the worker.
 func (h *Handler) RecordWorkerHeartbeat(
-	_ context.Context, request *matchingservice.RecordWorkerHeartbeatRequest,
+	ctx context.Context, request *matchingservice.RecordWorkerHeartbeatRequest,
 ) (*matchingservice.RecordWorkerHeartbeatResponse, error) {
 	nsID := namespace.ID(request.GetNamespaceId())
 	nsName := h.namespaceName(nsID)
+	principal := headers.GetPrincipal(ctx)
 
-	h.workersRegistry.RecordWorkerHeartbeats(nsID, nsName, request.GetHeartbeartRequest().GetWorkerHeartbeat())
+	h.workersRegistry.RecordWorkerHeartbeats(nsID, nsName, principal, request.GetHeartbeartRequest().GetWorkerHeartbeat())
 	return &matchingservice.RecordWorkerHeartbeatResponse{}, nil
 }
 

service/matching/workers/registry.go

@@ -1,6 +1,7 @@
 package workers
 
 import (
+	commonpb "go.temporal.io/api/common/v1"
 	workerpb "go.temporal.io/api/worker/v1"
 	"go.temporal.io/server/common/namespace"
 )
@@ -21,7 +22,7 @@ type (
 	}
 
 	Registry interface {
-		RecordWorkerHeartbeats(nsID namespace.ID, nsName namespace.Name, workerHeartbeat []*workerpb.WorkerHeartbeat)
+		RecordWorkerHeartbeats(nsID namespace.ID, nsName namespace.Name, principal *commonpb.Principal, workerHeartbeat []*workerpb.WorkerHeartbeat)
 		ListWorkers(nsID namespace.ID, params ListWorkersParams) (ListWorkersResponse, error)
 		DescribeWorker(nsID namespace.ID, workerInstanceKey string) (*workerpb.WorkerHeartbeat, error)
 	}

service/matching/workers/registry_impl.go

@@ -10,6 +10,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	commonpb "go.temporal.io/api/common/v1"
 	enumspb "go.temporal.io/api/enums/v1"
 	"go.temporal.io/api/serviceerror"
 	workerpb "go.temporal.io/api/worker/v1"
@@ -34,7 +35,7 @@ type (
 		hb             *workerpb.WorkerHeartbeat
 		lastSeen       time.Time
 		elem           *list.Element
-		isSystemWorkerWorker bool
+		isSystemWorker bool
 	}
 	// bucket holds part of the keyspace: a map from namespace → (map of instanceKey → entry),
 	// plus a recency list for eviction.
@@ -82,7 +83,7 @@ func newBucket() *bucket {
 // upsertHeartbeats inserts or refreshes a WorkerHeartbeat under the given namespace.
 // Returns the count of added and removed entries separately.
 // Workers with WORKER_STATUS_SHUTDOWN are immediately removed from the registry.
-func (b *bucket) upsertHeartbeats(nsID namespace.ID, heartbeats []*workerpb.WorkerHeartbeat) (added int64, removed int64) {
+func (b *bucket) upsertHeartbeats(nsID namespace.ID, principal *commonpb.Principal, heartbeats []*workerpb.WorkerHeartbeat) (added int64, removed int64) {
 	now := time.Now()
 
 	b.mu.Lock()
@@ -107,20 +108,20 @@ func (b *bucket) upsertHeartbeats(nsID namespace.ID, heartbeats []*workerpb.Work
 			continue
 		}
 
-		isSystemWorker := primitives.IsInternalTaskQueue(hb.GetTaskQueue())
+		isSystemWorker := isSystemPrincipal(principal, hb.GetTaskQueue())
 
 		// Normal upsert
 		if e, exists := mp[key]; exists {
 			e.hb = hb
 			e.lastSeen = now
-			e.isSystemWorkerWorker = isSystemWorker
+			e.isSystemWorker = isSystemWorker
 			b.order.MoveToBack(e.elem)
 		} else {
 			e = &entry{
 				nsID:           nsID,
 				hb:             hb,
 				lastSeen:       now,
-				isSystemWorkerWorker: isSystemWorker,
+				isSystemWorker: isSystemWorker,
 			}
 			e.elem = b.order.PushBack(e)
 			mp[key] = e
@@ -148,7 +149,7 @@ func (b *bucket) filterWorkers(
 	}
 	out := make([]*workerpb.WorkerHeartbeat, 0, len(mp))
 	for _, e := range mp {
-		if !includeSystemWorkers && e.isSystemWorkerWorker {
+		if !includeSystemWorkers && e.isSystemWorker {
 			continue
 		}
 		if predicate(e.hb) {
@@ -257,9 +258,9 @@ func (m *registryImpl) getBucket(nsID namespace.ID) *bucket {
 
 // upsertHeartbeat records or refreshes a WorkerHeartbeat under the given namespace.
 // New entries increment the global counter.
-func (m *registryImpl) upsertHeartbeats(nsID namespace.ID, heartbeats []*workerpb.WorkerHeartbeat) {
+func (m *registryImpl) upsertHeartbeats(nsID namespace.ID, principal *commonpb.Principal, heartbeats []*workerpb.WorkerHeartbeat) {
 	b := m.getBucket(nsID)
-	added, removed := b.upsertHeartbeats(nsID, heartbeats)
+	added, removed := b.upsertHeartbeats(nsID, principal, heartbeats)
 	m.total.Add(added - removed)
 	if added > 0 {
 		metrics.WorkerRegistryWorkersAdded.With(m.metricsHandler).Record(added)
@@ -375,8 +376,8 @@ func (m *registryImpl) Stop() {
 	close(m.quit)
 }
 
-func (m *registryImpl) RecordWorkerHeartbeats(nsID namespace.ID, nsName namespace.Name, workerHeartbeat []*workerpb.WorkerHeartbeat) {
-	m.upsertHeartbeats(nsID, workerHeartbeat)
+func (m *registryImpl) RecordWorkerHeartbeats(nsID namespace.ID, nsName namespace.Name, principal *commonpb.Principal, workerHeartbeat []*workerpb.WorkerHeartbeat) {
+	m.upsertHeartbeats(nsID, principal, workerHeartbeat)
 	m.metricsEmitter.emit(nsID, nsName, workerHeartbeat)
 }
 
@@ -477,3 +478,14 @@ func (m *registryImpl) DescribeWorker(nsID namespace.ID, workerInstanceKey strin
 	}
 	return b.getWorkerHeartbeat(nsID, workerInstanceKey)
 }
+
+// isSystemPrincipal determines if a worker is a system worker.
+// If a principal is available, it checks whether the principal type indicates
+// an internal server identity. Otherwise, it falls back to checking the task
+// queue name prefix.
+func isSystemPrincipal(principal *commonpb.Principal, taskQueue string) bool {
+	if principal != nil {
+		return principal.GetType() == "temporal"
+	}
+	return primitives.IsInternalTaskQueue(taskQueue)
+}

service/matching/workers/registry_impl_test.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	commonpb "go.temporal.io/api/common/v1"
 	enumspb "go.temporal.io/api/enums/v1"
 	workerpb "go.temporal.io/api/worker/v1"
 	"go.temporal.io/server/common/dynamicconfig"
@@ -47,7 +48,7 @@ func TestUpdateAndListNamespace(t *testing.T) {
 	// Add some heartbeats
 	hb1 := &workerpb.WorkerHeartbeat{WorkerInstanceKey: "workerA", Status: enumspb.WORKER_STATUS_RUNNING}
 	hb2 := &workerpb.WorkerHeartbeat{WorkerInstanceKey: "workerB", Status: enumspb.WORKER_STATUS_RUNNING}
-	m.upsertHeartbeats("ns1", []*workerpb.WorkerHeartbeat{hb1, hb2})
+	m.upsertHeartbeats("ns1", nil /*principal*/,[]*workerpb.WorkerHeartbeat{hb1, hb2})
 
 	list = m.filterWorkers("ns1", true /*includeSystemWorkers*/,alwaysTrue)
 	// Order is not guaranteed; check contents by keys
@@ -92,7 +93,7 @@ func TestShutdownStatusRemovesWorker(t *testing.T) {
 	// Add two running workers
 	hb1 := &workerpb.WorkerHeartbeat{WorkerInstanceKey: "worker1", Status: enumspb.WORKER_STATUS_RUNNING}
 	hb2 := &workerpb.WorkerHeartbeat{WorkerInstanceKey: "worker2", Status: enumspb.WORKER_STATUS_RUNNING}
-	m.upsertHeartbeats("ns1", []*workerpb.WorkerHeartbeat{hb1, hb2})
+	m.upsertHeartbeats("ns1", nil /*principal*/,[]*workerpb.WorkerHeartbeat{hb1, hb2})
 
 	// Verify both workers are registered
 	list := m.filterWorkers("ns1", true /*includeSystemWorkers*/,alwaysTrue)
@@ -101,7 +102,7 @@ func TestShutdownStatusRemovesWorker(t *testing.T) {
 
 	// Worker1 sends shutdown status
 	hbShutdown := &workerpb.WorkerHeartbeat{WorkerInstanceKey: "worker1", Status: enumspb.WORKER_STATUS_SHUTDOWN}
-	m.upsertHeartbeats("ns1", []*workerpb.WorkerHeartbeat{hbShutdown})
+	m.upsertHeartbeats("ns1", nil /*principal*/,[]*workerpb.WorkerHeartbeat{hbShutdown})
 
 	// Verify only worker1 is removed, worker2 remains
 	list = m.filterWorkers("ns1", true /*includeSystemWorkers*/,alwaysTrue)
@@ -139,7 +140,7 @@ func TestShutdownStatusForNonExistentWorker(t *testing.T) {
 
 	// Send shutdown for non-existent worker - should be a no-op
 	hb := &workerpb.WorkerHeartbeat{WorkerInstanceKey: "unknown", Status: enumspb.WORKER_STATUS_SHUTDOWN}
-	m.upsertHeartbeats("ns1", []*workerpb.WorkerHeartbeat{hb})
+	m.upsertHeartbeats("ns1", nil /*principal*/,[]*workerpb.WorkerHeartbeat{hb})
 
 	// Verify nothing happened
 	list := m.filterWorkers("ns1", true /*includeSystemWorkers*/,alwaysTrue)
@@ -165,7 +166,7 @@ func TestListNamespacePredicate(t *testing.T) {
 	for i := 1; i <= 5; i++ {
 		key := fmt.Sprintf("key%d", i)
 		hb := &workerpb.WorkerHeartbeat{WorkerInstanceKey: key, CurrentStickyCacheSize: int32(i)}
-		m.upsertHeartbeats("ns", []*workerpb.WorkerHeartbeat{hb})
+		m.upsertHeartbeats("ns", nil /*principal*/,[]*workerpb.WorkerHeartbeat{hb})
 	}
 
 	// Table-driven tests for predicates
@@ -191,7 +192,7 @@ func TestFilterWorkersExcludesSystemWorkers(t *testing.T) {
 	m := newRegistryImpl(testDefaultRegistryParams(metrics.NoopMetricsHandler))
 	defer m.Stop()
 
-	m.upsertHeartbeats("ns", []*workerpb.WorkerHeartbeat{
+	m.upsertHeartbeats("ns", nil /*principal*/,[]*workerpb.WorkerHeartbeat{
 		{WorkerInstanceKey: "user-1", TaskQueue: "my-queue"},
 		{WorkerInstanceKey: "user-2", TaskQueue: "my-queue"},
 		{WorkerInstanceKey: "sys-1", TaskQueue: "temporal-sys-per-ns-tq"},
@@ -211,6 +212,52 @@ func TestFilterWorkersExcludesSystemWorkers(t *testing.T) {
 	})
 }
 
+func TestIsSystemPrincipal(t *testing.T) {
+	t.Run("principal with type temporal marks as system worker", func(t *testing.T) {
+		m := newRegistryImpl(testDefaultRegistryParams(metrics.NoopMetricsHandler))
+		defer m.Stop()
+
+		principal := &commonpb.Principal{Type: "temporal", Name: "internal"}
+		m.upsertHeartbeats("ns", principal, []*workerpb.WorkerHeartbeat{
+			{WorkerInstanceKey: "sys-worker", TaskQueue: "any-queue"},
+		})
+
+		list := m.filterWorkers("ns", false /*includeSystemWorkers*/, alwaysTrue)
+		require.Empty(t, list, "worker with temporal principal should be excluded")
+
+		list = m.filterWorkers("ns", true /*includeSystemWorkers*/, alwaysTrue)
+		require.Len(t, list, 1)
+	})
+
+	t.Run("principal with non-temporal type is not system worker", func(t *testing.T) {
+		m := newRegistryImpl(testDefaultRegistryParams(metrics.NoopMetricsHandler))
+		defer m.Stop()
+
+		principal := &commonpb.Principal{Type: "user", Name: "alice"}
+		m.upsertHeartbeats("ns", principal, []*workerpb.WorkerHeartbeat{
+			{WorkerInstanceKey: "user-worker", TaskQueue: "temporal-sys-per-ns-tq"},
+		})
+
+		// Even though task queue has system prefix, principal says it's a user
+		list := m.filterWorkers("ns", false /*includeSystemWorkers*/, alwaysTrue)
+		require.Len(t, list, 1, "worker with user principal should not be excluded")
+	})
+
+	t.Run("nil principal falls back to task queue prefix", func(t *testing.T) {
+		m := newRegistryImpl(testDefaultRegistryParams(metrics.NoopMetricsHandler))
+		defer m.Stop()
+
+		m.upsertHeartbeats("ns", nil /*principal*/, []*workerpb.WorkerHeartbeat{
+			{WorkerInstanceKey: "sys-worker", TaskQueue: "temporal-sys-per-ns-tq"},
+			{WorkerInstanceKey: "user-worker", TaskQueue: "my-queue"},
+		})
+
+		list := m.filterWorkers("ns", false /*includeSystemWorkers*/, alwaysTrue)
+		require.Len(t, list, 1)
+		require.Equal(t, "user-worker", list[0].WorkerInstanceKey)
+	})
+}
+
 func TestEvictByTTL(t *testing.T) {
 	// Use capture handler to verify metrics
 	captureHandler := metricstest.NewCaptureHandler()
@@ -231,7 +278,7 @@ func TestEvictByTTL(t *testing.T) {
 	defer m.Stop()
 
 	hb := &workerpb.WorkerHeartbeat{WorkerInstanceKey: "oldWorker"}
-	m.upsertHeartbeats("ns", []*workerpb.WorkerHeartbeat{hb})
+	m.upsertHeartbeats("ns", nil /*principal*/,[]*workerpb.WorkerHeartbeat{hb})
 
 	// Manually move beyond TTL
 	b := m.getBucket("ns")
@@ -279,7 +326,7 @@ func TestEvictByCapacity(t *testing.T) {
 	for i := 1; i <= 5; i++ {
 		key := fmt.Sprintf("cap%d", i)
 		hb := &workerpb.WorkerHeartbeat{WorkerInstanceKey: key}
-		m.upsertHeartbeats("ns", []*workerpb.WorkerHeartbeat{hb})
+		m.upsertHeartbeats("ns", nil /*principal*/,[]*workerpb.WorkerHeartbeat{hb})
 	}
 
 	// All entries have lastSeen.Before(now) when MinEvictAge=0, so eligible
@@ -337,7 +384,7 @@ func TestEvictByCapacityWithMinAgeProtection(t *testing.T) {
 	for i := 1; i <= 3; i++ {
 		key := fmt.Sprintf("worker%d", i)
 		hb := &workerpb.WorkerHeartbeat{WorkerInstanceKey: key}
-		m.upsertHeartbeats("ns", []*workerpb.WorkerHeartbeat{hb})
+		m.upsertHeartbeats("ns", nil /*principal*/,[]*workerpb.WorkerHeartbeat{hb})
 	}
 
 	// Verify we're over capacity
@@ -389,7 +436,7 @@ func TestEvictByCapacityAfterMinAge(t *testing.T) {
 		for i := 1; i <= 3; i++ {
 			key := fmt.Sprintf("worker%d", i)
 			hb := &workerpb.WorkerHeartbeat{WorkerInstanceKey: key}
-			m.upsertHeartbeats("ns", []*workerpb.WorkerHeartbeat{hb})
+			m.upsertHeartbeats("ns", nil /*principal*/,[]*workerpb.WorkerHeartbeat{hb})
 		}
 
 		// Virtual time advance - instant with synctest!
@@ -439,14 +486,14 @@ func TestMultipleNamespaces(t *testing.T) {
 		{WorkerInstanceKey: "ns1-worker2", TaskQueue: "queue1"},
 		{WorkerInstanceKey: "ns1-worker3", TaskQueue: "queue2"},
 	}
-	m.upsertHeartbeats("namespace1", ns1Workers)
+	m.upsertHeartbeats("namespace1", nil /*principal*/,ns1Workers)
 
 	// Add 2 workers to namespace2
 	ns2Workers := []*workerpb.WorkerHeartbeat{
 		{WorkerInstanceKey: "ns2-worker1", TaskQueue: "queue3"},
 		{WorkerInstanceKey: "ns2-worker2", TaskQueue: "queue3"},
 	}
-	m.upsertHeartbeats("namespace2", ns2Workers)
+	m.upsertHeartbeats("namespace2", nil /*principal*/,ns2Workers)
 
 	// Verify functional behavior first
 	ns1List := m.filterWorkers("namespace1", true /*includeSystemWorkers*/,alwaysTrue)
@@ -504,7 +551,7 @@ func TestEvictLoopRecordsUtilizationMetric(t *testing.T) {
 		for i := 1; i <= 3; i++ {
 			key := fmt.Sprintf("worker%d", i)
 			hb := &workerpb.WorkerHeartbeat{WorkerInstanceKey: key}
-			m.upsertHeartbeats("ns", []*workerpb.WorkerHeartbeat{hb})
+			m.upsertHeartbeats("ns", nil /*principal*/,[]*workerpb.WorkerHeartbeat{hb})
 		}
 
 		// Verify initial state
@@ -548,7 +595,7 @@ func BenchmarkUpdate(b *testing.B) {
 	hb := &workerpb.WorkerHeartbeat{WorkerInstanceKey: "benchWorker"}
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		m.upsertHeartbeats("benchNs", []*workerpb.WorkerHeartbeat{hb})
+		m.upsertHeartbeats("benchNs", nil /*principal*/,[]*workerpb.WorkerHeartbeat{hb})
 	}
 }
 
@@ -569,7 +616,7 @@ func BenchmarkListNamespace(b *testing.B) {
 	for i := range 1000 {
 		key := fmt.Sprintf("worker%d", i)
 		hb := &workerpb.WorkerHeartbeat{WorkerInstanceKey: key}
-		m.upsertHeartbeats("benchNs", []*workerpb.WorkerHeartbeat{hb})
+		m.upsertHeartbeats("benchNs", nil /*principal*/,[]*workerpb.WorkerHeartbeat{hb})
 	}
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
@@ -604,7 +651,7 @@ func BenchmarkRandomUpdate(b *testing.B) {
 		for i := range totalHeartbeats {
 			key := fmt.Sprintf("%s-worker%d", ns, i)
 			hb := &workerpb.WorkerHeartbeat{WorkerInstanceKey: key, CurrentStickyCacheSize: int32(i)}
-			m.upsertHeartbeats(ns, []*workerpb.WorkerHeartbeat{hb})
+			m.upsertHeartbeats(ns, nil /*principal*/, []*workerpb.WorkerHeartbeat{hb})
 			pairs = append(pairs, pair{ns: ns, hb: hb})
 		}
 	}
@@ -613,7 +660,7 @@ func BenchmarkRandomUpdate(b *testing.B) {
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
 		p := pairs[r.Intn(len(pairs))]
-		m.upsertHeartbeats(p.ns, []*workerpb.WorkerHeartbeat{p.hb})
+		m.upsertHeartbeats(p.ns, nil /*principal*/, []*workerpb.WorkerHeartbeat{p.hb})
 	}
 }
 
@@ -662,7 +709,7 @@ func TestActivitySlotsMetric(t *testing.T) {
 
 	testNamespace := tv.NamespaceID()
 	testNamespaceName := namespace.Name(testNamespace + "_name")
-	m.RecordWorkerHeartbeats(testNamespace, testNamespaceName, []*workerpb.WorkerHeartbeat{worker1, worker2, worker3})
+	m.RecordWorkerHeartbeats(testNamespace, testNamespaceName, nil /*principal*/, []*workerpb.WorkerHeartbeat{worker1, worker2, worker3})
 
 	// Verify activity slots metrics
 	snapshot := capture.Snapshot()
@@ -739,7 +786,7 @@ func TestPluginMetricsExported(t *testing.T) {
 
 	testNamespace := tv.NamespaceID()
 	testNamespaceName := namespace.Name(testNamespace + "_name")
-	m.RecordWorkerHeartbeats(testNamespace, testNamespaceName, []*workerpb.WorkerHeartbeat{worker1, worker2, worker3})
+	m.RecordWorkerHeartbeats(testNamespace, testNamespaceName, nil /*principal*/, []*workerpb.WorkerHeartbeat{worker1, worker2, worker3})
 
 	// Verify plugin metrics - should have exactly 3 recordings despite plugin-a being in both workers
 	snapshot := capture.Snapshot()
@@ -801,7 +848,7 @@ func TestPluginMetricsDisabled(t *testing.T) {
 	// Upsert heartbeats for test namespace
 	testNamespace := tv.NamespaceID()
 	testNamespaceName := namespace.Name(testNamespace + "_name")
-	m.RecordWorkerHeartbeats(testNamespace, testNamespaceName, []*workerpb.WorkerHeartbeat{worker1})
+	m.RecordWorkerHeartbeats(testNamespace, testNamespaceName, nil /*principal*/, []*workerpb.WorkerHeartbeat{worker1})
 
 	// Verify no plugin metrics are recorded when disabled
 	snapshot := capture.Snapshot()