fix(authz): reuse policy and entity resolution within bulk requests

service/authorization/v2/authorization.go

@@ -277,6 +277,8 @@ func (as *Service) GetDecisionBulk(ctx context.Context, req *connect.Request[aut
 		return nil, statusifyError(ctx, as.logger, errors.Join(ErrFailedToInitPDP, err))
 	}
 
+	pdp = pdp.WithRequestReuse()
+
 	multiRequests := req.Msg.GetDecisionRequests()
 	decisionResponses := make([]*authzV2.GetDecisionMultiResourceResponse, len(multiRequests))
 
@@ -285,8 +287,6 @@ func (as *Service) GetDecisionBulk(ctx context.Context, req *connect.Request[aut
 		return nil, statusifyError(ctx, as.logger, err)
 	}
 
-	// TODO: revisit performance of this loop after introduction of caching and registered resource values within decisioning,
-	// as the same entity in multiple requests should only be resolved JIT once, not once per request if the same in each.
 	for idx, request := range multiRequests {
 		entityIdentifier := request.GetEntityIdentifier()
 		action := request.GetAction()

service/internal/access/v2/just_in_time_pdp.go

@@ -38,6 +38,7 @@ var (
 )
 
 type JustInTimePDP struct {
+	reuse  *requestReuse
 	logger *logger.Logger
 	sdk    *otdfSDK.SDK
 	// embedded obligations PDP
@@ -346,7 +347,7 @@ func (p *JustInTimePDP) GetEntitlements(
 // buildInnerPDP fetches the entitleable attributes for the provided value FQNs and constructs a
 // request-scoped PolicyDecisionPoint from them plus the fully-loaded registered resources and
 // dynamic value mappings.
-func (p *JustInTimePDP) buildInnerPDP(ctx context.Context, valueFQNs []string) (*PolicyDecisionPoint, error) {
+func (p *JustInTimePDP) buildInnerPDPUncached(ctx context.Context, valueFQNs []string) (*PolicyDecisionPoint, error) {
 	// Direct entitlements / dynamic value mappings require the full policy load (see NewJustInTimePDP).
 	if p.fullPolicyPDP != nil {
 		return p.fullPolicyPDP, nil
@@ -457,7 +458,7 @@ func (p *JustInTimePDP) getMatchedSubjectMappings(
 }
 
 // resolveEntitiesFromEntityChain roundtrips caller-provided entity chains through ERS.
-func (p *JustInTimePDP) resolveEntitiesFromEntityChain(
+func (p *JustInTimePDP) resolveEntitiesFromEntityChainUncached(
 	ctx context.Context,
 	entityChain *entity.EntityChain,
 	skipEnvironmentEntities bool,
@@ -548,7 +549,7 @@ func entityRepresentationsFromResolvedChain(entityChain *entity.EntityChain, ski
 
 // resolveEntitiesFromToken roundtrips to ERS to resolve the provided token
 // and optionally skips environment entities (which is expected behavior in decision flow)
-func (p *JustInTimePDP) resolveEntitiesFromToken(
+func (p *JustInTimePDP) resolveEntitiesFromTokenUncached(
 	ctx context.Context,
 	token *entity.Token,
 	skipEnvironmentEntities bool,

service/internal/access/v2/request_reuse.go

@@ -0,0 +1,95 @@
+package access
+
+import (
+	"context"
+	"fmt"
+	"slices"
+	"strings"
+
+	authzV2 "github.com/opentdf/platform/protocol/go/authorization/v2"
+	"github.com/opentdf/platform/protocol/go/entity"
+	entityresolutionV2 "github.com/opentdf/platform/protocol/go/entityresolution/v2"
+	attrs "github.com/opentdf/platform/protocol/go/policy/attributes"
+	"google.golang.org/protobuf/proto"
+)
+
+const (
+	maxRequestPolicyCacheEntries     = 8
+	maxRequestResolutionCacheEntries = 32
+)
+
+type resolutionCacheKey struct {
+	messageType     string
+	request         string
+	skipEnvironment bool
+}
+
+type requestReuse struct {
+	policies    map[string]*PolicyDecisionPoint
+	resolutions map[resolutionCacheKey][]*entityresolutionV2.EntityRepresentation
+}
+
+// WithRequestReuse returns a PDP for sequential subrequests within one bulk RPC.
+// Its bounded memo tables must never be shared across RPCs or caller contexts.
+func (p *JustInTimePDP) WithRequestReuse() *JustInTimePDP {
+	scoped := *p
+	scoped.reuse = &requestReuse{policies: make(map[string]*PolicyDecisionPoint), resolutions: make(map[resolutionCacheKey][]*entityresolutionV2.EntityRepresentation)}
+	return &scoped
+}
+
+func (p *JustInTimePDP) buildInnerPDP(ctx context.Context, fqns []string) (*PolicyDecisionPoint, error) {
+	if p.reuse == nil || p.fullPolicyPDP != nil {
+		return p.buildInnerPDPUncached(ctx, fqns)
+	}
+	normalized := make([]string, len(fqns))
+	for i, fqn := range fqns {
+		normalized[i] = strings.ToLower(fqn)
+	}
+	slices.Sort(normalized)
+	normalized = slices.Compact(normalized)
+	encoded, err := proto.Marshal(&attrs.GetEntitleableAttributesByFqnsRequest{Fqns: normalized})
+	if err != nil {
+		return nil, fmt.Errorf("encode policy lookup: %w", err)
+	}
+	key := string(encoded)
+	if cached, ok := p.reuse.policies[key]; ok {
+		return cached, nil
+	}
+	pdp, err := p.buildInnerPDPUncached(ctx, normalized)
+	if err == nil && len(p.reuse.policies) < maxRequestPolicyCacheEntries {
+		p.reuse.policies[key] = pdp
+	}
+	return pdp, err
+}
+
+func (p *JustInTimePDP) reuseResolution(request proto.Message, skipEnvironment bool, resolve func() ([]*entityresolutionV2.EntityRepresentation, error)) ([]*entityresolutionV2.EntityRepresentation, error) {
+	if p.reuse == nil {
+		return resolve()
+	}
+	encoded, err := proto.MarshalOptions{Deterministic: true}.Marshal(request)
+	if err != nil {
+		return nil, fmt.Errorf("encode entity resolution request: %w", err)
+	}
+	key := resolutionCacheKey{messageType: string(request.ProtoReflect().Descriptor().FullName()), request: string(encoded), skipEnvironment: skipEnvironment}
+	if cached, ok := p.reuse.resolutions[key]; ok {
+		return cached, nil
+	}
+	representations, err := resolve()
+	if err == nil && len(p.reuse.resolutions) < maxRequestResolutionCacheEntries {
+		p.reuse.resolutions[key] = representations
+	}
+	return representations, err
+}
+
+func (p *JustInTimePDP) resolveEntitiesFromEntityChain(ctx context.Context, chain *entity.EntityChain, skipEnvironment bool) ([]*entityresolutionV2.EntityRepresentation, error) {
+	return p.reuseResolution(chain, skipEnvironment, func() ([]*entityresolutionV2.EntityRepresentation, error) {
+		return p.resolveEntitiesFromEntityChainUncached(ctx, chain, skipEnvironment)
+	})
+}
+
+func (p *JustInTimePDP) resolveEntitiesFromToken(ctx context.Context, token *entity.Token, skipEnvironment bool, resources []*authzV2.Resource) ([]*entityresolutionV2.EntityRepresentation, error) {
+	request := &entityresolutionV2.CreateEntityChainsFromTokensRequest{Tokens: []*entity.Token{token}, Resources: resources}
+	return p.reuseResolution(request, skipEnvironment, func() ([]*entityresolutionV2.EntityRepresentation, error) {
+		return p.resolveEntitiesFromTokenUncached(ctx, token, skipEnvironment, resources)
+	})
+}

service/internal/access/v2/request_reuse_test.go

@@ -0,0 +1,72 @@
+package access
+
+import (
+	"fmt"
+	"testing"
+
+	authzV2 "github.com/opentdf/platform/protocol/go/authorization/v2"
+	"github.com/opentdf/platform/protocol/go/entity"
+	entityresolutionV2 "github.com/opentdf/platform/protocol/go/entityresolution/v2"
+	otdfSDK "github.com/opentdf/platform/sdk"
+	"github.com/opentdf/platform/service/logger"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/proto"
+)
+
+func TestRequestReuseCachesPolicySetsAndStaysRequestScoped(t *testing.T) {
+	const def = "https://example.com/attr/department"
+	const fqn = def + "/value/engineering"
+	client := decisionAttrFake(def, fqn, "abc")
+	original := &JustInTimePDP{logger: logger.CreateTestLogger(), sdk: &otdfSDK.SDK{Attributes: client}}
+	scoped := original.WithRequestReuse()
+	first, err := scoped.buildInnerPDP(t.Context(), []string{fqn})
+	require.NoError(t, err)
+	second, err := scoped.buildInnerPDP(t.Context(), []string{fqn, fqn})
+	require.NoError(t, err)
+	require.Same(t, first, second)
+	require.Len(t, client.requests, 1)
+	fresh := original.WithRequestReuse()
+	third, err := fresh.buildInnerPDP(t.Context(), []string{fqn})
+	require.NoError(t, err)
+	require.NotSame(t, first, third)
+	require.Len(t, client.requests, 2)
+	for i := 0; i < maxRequestPolicyCacheEntries+1; i++ {
+		_, err := scoped.buildInnerPDP(t.Context(), []string{fmt.Sprintf("%s/value/%d", def, i)})
+		require.NoError(t, err)
+	}
+	require.Len(t, scoped.reuse.policies, maxRequestPolicyCacheEntries)
+}
+
+func TestRequestReuseIncludesTokenResourcesAndEnvironmentMode(t *testing.T) {
+	client := &recordingERSV2Client{createResponse: &entityresolutionV2.CreateEntityChainsFromTokensResponse{
+		EntityChains: []*entity.EntityChain{{Entities: []*entity.Entity{{EphemeralId: "subject", Category: entity.Entity_CATEGORY_SUBJECT, EntityType: &entity.Entity_Claims{Claims: claimsAnyForTest(t, map[string]interface{}{"name": "alice"})}}}}},
+	}}
+	scoped := testJITPDP(client).WithRequestReuse()
+	token := &entity.Token{Jwt: "token", EphemeralId: "token-id"}
+	resources := attrValueResource("https://example.com/attr/a/value/one")
+	_, err := scoped.resolveEntitiesFromToken(t.Context(), token, true, resources)
+	require.NoError(t, err)
+	_, err = scoped.resolveEntitiesFromToken(t.Context(), proto.CloneOf(token), true, []*authzV2.Resource{proto.CloneOf(resources[0])})
+	require.NoError(t, err)
+	require.Equal(t, 1, client.createCalls)
+	_, err = scoped.resolveEntitiesFromToken(t.Context(), token, true, attrValueResource("https://example.com/attr/a/value/two"))
+	require.NoError(t, err)
+	require.Equal(t, 2, client.createCalls)
+	_, err = scoped.resolveEntitiesFromToken(t.Context(), token, false, resources)
+	require.NoError(t, err)
+	require.Equal(t, 3, client.createCalls)
+	_, err = scoped.resolveEntitiesFromToken(t.Context(), &entity.Token{Jwt: "different", EphemeralId: "token-id"}, true, resources)
+	require.NoError(t, err)
+	require.Equal(t, 4, client.createCalls)
+}
+
+func TestRequestReuseResolvesIdenticalEntityChainOnce(t *testing.T) {
+	client := &recordingERSV2Client{resolveResponse: &entityresolutionV2.ResolveEntitiesResponse{EntityRepresentations: []*entityresolutionV2.EntityRepresentation{entityRepWithClientID("abc")}}}
+	scoped := testJITPDP(client).WithRequestReuse()
+	chain := entityChainIdentifier().GetEntityChain()
+	for range 3 {
+		_, err := scoped.resolveEntitiesFromEntityChain(t.Context(), chain, true)
+		require.NoError(t, err)
+	}
+	require.Equal(t, 1, client.resolveCalls)
+}