Skip to content

perf(graphrag): bound relation authorization hot paths #321

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kl3inIT merged 1 commit into from
Aug 6, 2026
Merged

perf(graphrag): bound relation authorization hot paths #321

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions .../graph-rag-core/src/main/java/com/orgmemory/graphrag/observability/GraphRagEventSink.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -267,6 +267,17 @@ enum Stage {
AUTHORIZE,
PREPARE_QUERY,
RETRIEVE_SNAPSHOT,
QUERY_SEARCH_ENTITIES,
QUERY_SEARCH_RELATIONS,
QUERY_SEARCH_CHUNKS,
QUERY_EXPAND_ENTITY_IDS,
QUERY_LOAD_INCIDENT_RELATIONS,
QUERY_LOAD_ENTITY_CONTRIBUTIONS,
QUERY_LOAD_RELATION_CONTRIBUTIONS,
QUERY_LOAD_VISIBLE_ENTITY_DEGREES,
QUERY_LOAD_VISIBLE_RELATION_WEIGHTS,
QUERY_RANK_CHUNKS,
QUERY_LOAD_CHUNKS,
RETRIEVE,
RERANK,
ASSEMBLE_CONTEXT
Expand Down
248 changes: 201 additions & 47 deletions ...onents/graph-rag-core/src/main/java/com/orgmemory/graphrag/query/LightRagQueryEngine.java
View file
Open in desktop

Large diffs are not rendered by default.

32 changes: 32 additions & 0 deletions ...it/src/test/java/com/orgmemory/graphrag/testkit/LightRagQueryRuntimeConformanceTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,38 @@ void mixBatchesRawQueryLowKeywordsAndHighKeywordsExactlyOnce() {
.anyMatch(signal -> signal.origin() == LightRagQueryResult.Origin.RELATION));
}

@Test
void preparedExecutionReportsBoundedProjectionOperationsWithoutPayload() {
LightRagQueryRequest request = request(
LightRagQueryMode.MIX,
QueryOutputMode.CONTEXT,
false,
true,
false,
trustedKeywords());
var prepared = engine.prepare(request);
List<LightRagQueryEngine.QueryOperationMeasurement> measurements =
new ArrayList<>();

LightRagQueryResult result =
engine.executePrepared(request, prepared, measurements::add);

assertEquals(LightRagQueryResult.Status.SUCCESS, result.status());
assertTrue(measurements.stream().anyMatch(measurement ->
measurement.operation()
== LightRagQueryEngine.QueryOperation.SEARCH_ENTITIES));
assertTrue(measurements.stream().anyMatch(measurement ->
measurement.operation()
== LightRagQueryEngine.QueryOperation.EXPAND_ENTITY_IDS));
assertTrue(measurements.stream().anyMatch(measurement ->
measurement.operation()
== LightRagQueryEngine.QueryOperation.LOAD_INCIDENT_RELATIONS));
assertTrue(measurements.stream().allMatch(measurement ->
!measurement.duration().isNegative()
&& measurement.inputCount() >= 0
&& measurement.outputCount() >= 0));
}

@Test
void onePreparedQueryReusesKeywordAndEmbeddingEffectsAcrossSnapshotExecutions() {
LightRagQueryRequest request = request(
Expand Down
65 changes: 61 additions & 4 deletions ...java/com/orgmemory/core/knowledge/retrieval/DefaultGraphRagKnowledgeRetrievalService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -636,7 +636,10 @@ private PublishedSpaceQuery queryPublishedSpaces(
futures.add(completed.submit(tasks.decorate(() ->
new IndexedSnapshotQueryResult(
resultIndex,
queryPublishedSpace(request, prepared)))));
queryPublishedSpace(
operationId,
request,
prepared)))));
}
SnapshotQueryResult[] ordered =
new SnapshotQueryResult[requests.size()];
Expand Down Expand Up @@ -679,12 +682,18 @@ private PublishedSpaceQuery queryPublishedSpaces(
}

private SnapshotQueryResult queryPublishedSpace(
UUID operationId,
LightRagQueryRequest request,
LightRagPreparedQuery prepared) throws Exception {
long startedAt = System.nanoTime();
LightRagQueryResult result =
admission.execute(() ->
engine.executePrepared(request, prepared));
LightRagQueryResult result = admission.execute(() -> engine.executePrepared(
request,
prepared,
measurement -> emitQueryOperation(
operationId,
request.scope().organizationId(),
request.scope().authorizationFingerprint(),
measurement)));
return new SnapshotQueryResult(
result,
Duration.ofNanos(Math.max(
Expand All @@ -694,6 +703,54 @@ private SnapshotQueryResult queryPublishedSpace(
request.snapshot().namespace());
}

private void emitQueryOperation(
UUID operationId,
UUID organizationId,
String scopeFingerprint,
LightRagQueryEngine.QueryOperationMeasurement measurement) {
GraphRagEventSink.Outcome outcome = switch (measurement.outcome()) {
case SUCCEEDED -> GraphRagEventSink.Outcome.SUCCEEDED;
case FAILED -> GraphRagEventSink.Outcome.FAILED;
};
safeEmit(new GraphRagEventSink.GraphRagEvent(
operationId,
organizationId,
queryOperationStage(measurement.operation()),
outcome,
measurement.duration(),
measurement.inputCount(),
measurement.outputCount(),
null,
scopeFingerprint,
null,
outcome == GraphRagEventSink.Outcome.FAILED
? "query_operation_failed"
: null,
Instant.now()));
}

private static GraphRagEventSink.Stage queryOperationStage(
LightRagQueryEngine.QueryOperation operation) {
return switch (operation) {
case SEARCH_ENTITIES -> GraphRagEventSink.Stage.QUERY_SEARCH_ENTITIES;
case SEARCH_RELATIONS -> GraphRagEventSink.Stage.QUERY_SEARCH_RELATIONS;
case SEARCH_CHUNKS -> GraphRagEventSink.Stage.QUERY_SEARCH_CHUNKS;
case EXPAND_ENTITY_IDS -> GraphRagEventSink.Stage.QUERY_EXPAND_ENTITY_IDS;
case LOAD_INCIDENT_RELATIONS ->
GraphRagEventSink.Stage.QUERY_LOAD_INCIDENT_RELATIONS;
case LOAD_ENTITY_CONTRIBUTIONS ->
GraphRagEventSink.Stage.QUERY_LOAD_ENTITY_CONTRIBUTIONS;
case LOAD_RELATION_CONTRIBUTIONS ->
GraphRagEventSink.Stage.QUERY_LOAD_RELATION_CONTRIBUTIONS;
case LOAD_VISIBLE_ENTITY_DEGREES ->
GraphRagEventSink.Stage.QUERY_LOAD_VISIBLE_ENTITY_DEGREES;
case LOAD_VISIBLE_RELATION_WEIGHTS ->
GraphRagEventSink.Stage.QUERY_LOAD_VISIBLE_RELATION_WEIGHTS;
case RANK_CHUNKS -> GraphRagEventSink.Stage.QUERY_RANK_CHUNKS;
case LOAD_CHUNKS -> GraphRagEventSink.Stage.QUERY_LOAD_CHUNKS;
};
}

private record IndexedSnapshotQueryResult(
int index,
SnapshotQueryResult result) {
Expand Down
67 changes: 48 additions & 19 deletions ...t/java/com/orgmemory/core/knowledge/retrieval/GraphRagKnowledgeRetrievalServiceTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,7 +125,7 @@ void multipleSpacesPrepareOneLogicalQueryBeforeSnapshotRetrieval() {
preparedQueryPlan();
when(engine.prepare(any())).thenReturn(prepared);
CountDownLatch bothSpacesStarted = new CountDownLatch(2);
when(engine.executePrepared(any(), any())).thenAnswer(invocation -> {
when(engine.executePrepared(any(), any(), any())).thenAnswer(invocation -> {
bothSpacesStarted.countDown();
assertTrue(
bothSpacesStarted.await(2, TimeUnit.SECONDS),
Expand All @@ -152,7 +152,7 @@ void multipleSpacesPrepareOneLogicalQueryBeforeSnapshotRetrieval() {
assertEquals(List.of(), result.evidence());
verify(engine).prepare(any());
verify(engine, times(2))
.executePrepared(any(), any());
.executePrepared(any(), any(), any());
ArgumentCaptor<GraphRagEventSink.GraphRagEvent> captured =
ArgumentCaptor.forClass(GraphRagEventSink.GraphRagEvent.class);
verify(events, atLeastOnce()).emit(captured.capture());
Expand Down Expand Up @@ -190,7 +190,7 @@ void acquiresAdmissionPermitBeforeExecutingTheSnapshotStoreQuery() {
LightRagQueryEngine engine = mock(LightRagQueryEngine.class);
LightRagPreparedQuery prepared = preparedQueryPlan();
when(engine.prepare(any())).thenReturn(prepared);
when(engine.executePrepared(any(), any())).thenAnswer(invocation -> {
when(engine.executePrepared(any(), any(), any())).thenAnswer(invocation -> {
assertEquals(
0,
admission.availablePermits(),
Expand All @@ -213,7 +213,7 @@ void acquiresAdmissionPermitBeforeExecutingTheSnapshotStoreQuery() {
10,
"request-admission-before-checkout");

verify(engine).executePrepared(any(), any());
verify(engine).executePrepared(any(), any(), any());
assertEquals(1, admission.availablePermits());
}

Expand All @@ -232,7 +232,7 @@ void continuousAdmissionDoesNotWaitForAnEarlierSnapshotBatch()
CountDownLatch firstStarted = new CountDownLatch(1);
CountDownLatch releaseFirst = new CountDownLatch(1);
CountDownLatch thirdStarted = new CountDownLatch(1);
when(engine.executePrepared(any(), any())).thenAnswer(invocation -> {
when(engine.executePrepared(any(), any(), any())).thenAnswer(invocation -> {
LightRagQueryRequest request = invocation.getArgument(0);
Set<UUID> assets = request.scope().authorizedAssetIds();
if (assets.contains(ASSET_ID)) {
Expand Down Expand Up @@ -285,7 +285,7 @@ void snapshotFailureCancelsOutstandingContinuouslyAdmittedTasks()
CountDownLatch blockersStarted = new CountDownLatch(2);
CountDownLatch neverReleased = new CountDownLatch(1);
CountDownLatch interrupted = new CountDownLatch(2);
when(engine.executePrepared(any(), any())).thenAnswer(invocation -> {
when(engine.executePrepared(any(), any(), any())).thenAnswer(invocation -> {
LightRagQueryRequest request = invocation.getArgument(0);
if (request.scope().authorizedAssetIds().contains(SECOND_ASSET_ID)) {
assertTrue(blockersStarted.await(2, TimeUnit.SECONDS));
Expand Down Expand Up @@ -351,11 +351,20 @@ void retrievalObservationRunsKeywordAndRawQueryPathsWithoutAnswerGeneration() {
? keywordPrepared
: bypassPrepared;
});
when(engine.executePrepared(any(), any())).thenReturn(queryResult(
allowed.forKnowledgeSpace(SPACE_ID).authorizationFingerprint(),
grounding,
false,
false));
when(engine.executePrepared(any(), any(), any())).thenAnswer(invocation -> {
LightRagQueryEngine.QueryOperationObserver observer = invocation.getArgument(2);
observer.observe(new LightRagQueryEngine.QueryOperationMeasurement(
LightRagQueryEngine.QueryOperation.LOAD_INCIDENT_RELATIONS,
LightRagQueryEngine.QueryOperationOutcome.SUCCEEDED,
Duration.ofMillis(17),
12,
9));
return queryResult(
allowed.forKnowledgeSpace(SPACE_ID).authorizationFingerprint(),
grounding,
false,
false);
});
when(engine.consolidateGrounding(any(), any(), any())).thenReturn(preparedGrounding);
when(engine.renderGrounding(any(), any(), any())).thenReturn(preparedGrounding);
RelationshipAuthorizationSetPort finalAuthorization =
Expand All @@ -369,14 +378,15 @@ void retrievalObservationRunsKeywordAndRawQueryPathsWithoutAnswerGeneration() {
candidate(ENTITY_CHUNK_ID),
candidate(RELATION_CHUNK_ID),
candidate(CHUNK_ID)));
GraphRagEventSink events = mock(GraphRagEventSink.class);
GraphRagKnowledgeRetrievalService service = service(
scopes,
finalAuthorization,
canonical,
engine,
GraphRagRetrievalPolicy.defaults(),
audit,
mock(GraphRagEventSink.class));
events);

GraphRagKnowledgeRetrievalService.RetrievalObservation observation = service.observe(
actor,
Expand Down Expand Up @@ -404,6 +414,25 @@ void retrievalObservationRunsKeywordAndRawQueryPathsWithoutAnswerGeneration() {
assertEquals("model", observation.keywordPlan().source());
assertEquals(3, observation.keywordSeededDocuments().size());
assertEquals(3, observation.bypassDocuments().size());
ArgumentCaptor<GraphRagEventSink.GraphRagEvent> emitted =
ArgumentCaptor.forClass(GraphRagEventSink.GraphRagEvent.class);
verify(events, atLeastOnce()).emit(emitted.capture());
GraphRagEventSink.GraphRagEvent operation = emitted.getAllValues().stream()
.filter(event -> event.stage()
== GraphRagEventSink.Stage.QUERY_LOAD_INCIDENT_RELATIONS)
.findFirst()
.orElseThrow();
GraphRagEventSink.GraphRagEvent retrieval = emitted.getAllValues().stream()
.filter(event -> event.stage() == GraphRagEventSink.Stage.RETRIEVE)
.findFirst()
.orElseThrow();
assertEquals(Duration.ofMillis(17), operation.duration());
assertEquals(12, operation.inputCount());
assertEquals(9, operation.outputCount());
assertEquals(
allowed.forKnowledgeSpace(SPACE_ID).authorizationFingerprint(),
operation.scopeFingerprint());
assertEquals(retrieval.operationId(), operation.operationId());
}

@Test
Expand Down Expand Up @@ -483,7 +512,7 @@ void multipleSpacesDoNotInvokePerSpaceReranking() {
"request-multi-space-rerank"));

verify(engine, never()).prepare(any());
verify(engine, never()).executePrepared(any(), any());
verify(engine, never()).executePrepared(any(), any(), any());
}

@Test
Expand Down Expand Up @@ -531,7 +560,7 @@ void revocationBetweenRetrievalAndCitationCausesAFullRetryWithoutEgress() {
LightRagGrounding grounding = grounding();
LightRagGroundingAssembler.PreparedGrounding prepared =
prepared(grounding);
when(engine.executePrepared(any(), any())).thenReturn(queryResult(
when(engine.executePrepared(any(), any(), any())).thenReturn(queryResult(
allowed.forKnowledgeSpace(SPACE_ID)
.authorizationFingerprint(),
grounding,
Expand Down Expand Up @@ -585,7 +614,7 @@ void revocationBetweenRetrievalAndCitationCausesAFullRetryWithoutEgress() {

assertEquals(List.of(), result.evidence());
verify(engine).prepare(any());
verify(engine).executePrepared(any(), any());
verify(engine).executePrepared(any(), any(), any());
verify(finalAuthorization, never()).batchCheck(any());
assertEquals(0, canonical.recheckCount);
ArgumentCaptor<GraphRagEventSink.GraphRagEvent> captured =
Expand Down Expand Up @@ -624,7 +653,7 @@ void verifiesTheCompleteGraphGroundingBeforeCreatingTheModelInput() {
LightRagPreparedQuery queryPlan = preparedQueryPlan();
when(engine.prepare(any()))
.thenReturn(queryPlan);
when(engine.executePrepared(any(), any())).thenReturn(queryResult(
when(engine.executePrepared(any(), any(), any())).thenReturn(queryResult(
allowed.forKnowledgeSpace(SPACE_ID)
.authorizationFingerprint(),
grounding,
Expand Down Expand Up @@ -730,7 +759,7 @@ void contextAssemblyReportsWhatTheAnswerCostAndWhatTheBudgetRefused() {
LightRagQueryEngine engine = mock(LightRagQueryEngine.class);
LightRagPreparedQuery queryPlan = preparedQueryPlan();
when(engine.prepare(any())).thenReturn(queryPlan);
when(engine.executePrepared(any(), any())).thenReturn(queryResult(
when(engine.executePrepared(any(), any(), any())).thenReturn(queryResult(
allowed.forKnowledgeSpace(SPACE_ID).authorizationFingerprint(),
grounding,
true,
Expand Down Expand Up @@ -809,7 +838,7 @@ void authorizationModelMismatchCannotReachTheVerifiedRenderer() {
LightRagPreparedQuery queryPlan = preparedQueryPlan();
when(engine.prepare(any()))
.thenReturn(queryPlan);
when(engine.executePrepared(any(), any())).thenReturn(queryResult(
when(engine.executePrepared(any(), any(), any())).thenReturn(queryResult(
allowed.forKnowledgeSpace(SPACE_ID)
.authorizationFingerprint(),
grounding,
Expand Down Expand Up @@ -1002,7 +1031,7 @@ private static AuthorizationFixture authorizationFixture(
LightRagQueryEngine engine = mock(LightRagQueryEngine.class);
LightRagPreparedQuery queryPlan = preparedQueryPlan();
when(engine.prepare(any())).thenReturn(queryPlan);
when(engine.executePrepared(any(), any())).thenReturn(queryResult(
when(engine.executePrepared(any(), any(), any())).thenReturn(queryResult(
allowed.forKnowledgeSpace(SPACE_ID)
.authorizationFingerprint(),
grounding,
Expand Down
Loading
Loading