Skip to content

Commit

Permalink
[#4973] improvement(client-java/server): Clear CallerContext in Jav…
Browse files Browse the repository at this point in the history
…a Client and Server when call `getFileLocation` finally (#4974)

### What changes were proposed in this pull request?

Currently, after the client completes the `getFileLocation` request and
the server completes the `getFileLocation` response, the Thread Local
`CallerContext` is not cleaned up. This may cause the next request to
incorrectly report the same information repeatedly when the thread is
reused. We should eventually clean up the `CallerContext`.

### Why are the changes needed?

Fix: #4973 

### How was this patch tested?

Add UTs.
  • Loading branch information
xloya authored Sep 20, 2024
1 parent f54bfc1 commit 8aee4a3
Show file tree
Hide file tree
Showing 4 changed files with 36 additions and 14 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -244,20 +244,25 @@ public String getFileLocation(NameIdentifier ident, String subPath)
checkFilesetNameIdentifier(ident);
Namespace fullNamespace = getFilesetFullNamespace(ident.namespace());

CallerContext callerContext = CallerContext.CallerContextHolder.get();

Map<String, String> params = new HashMap<>();
params.put("sub_path", RESTUtils.encodeString(subPath));
FileLocationResponse resp =
restClient.get(
formatFileLocationRequestPath(fullNamespace, ident.name()),
params,
FileLocationResponse.class,
callerContext != null ? callerContext.context() : Collections.emptyMap(),
ErrorHandlers.filesetErrorHandler());
resp.validate();

return resp.getFileLocation();
try {
CallerContext callerContext = CallerContext.CallerContextHolder.get();

Map<String, String> params = new HashMap<>();
params.put("sub_path", RESTUtils.encodeString(subPath));
FileLocationResponse resp =
restClient.get(
formatFileLocationRequestPath(fullNamespace, ident.name()),
params,
FileLocationResponse.class,
callerContext != null ? callerContext.context() : Collections.emptyMap(),
ErrorHandlers.filesetErrorHandler());
resp.validate();

return resp.getFileLocation();
} finally {
// Clear the caller context
CallerContext.CallerContextHolder.remove();
}
}

@VisibleForTesting
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -534,6 +534,8 @@ public void testCallerContextToHeader() throws JsonProcessingException {
NameIdentifier.of(fileset.namespace().level(2), fileset.name()), mockSubPath);
Assertions.assertEquals(FilesetDataOperation.GET_FILE_STATUS.name(), dataOperation.get());
Assertions.assertEquals(InternalClientType.HADOOP_GVFS.name(), internalClientType.get());
// the caller context should be cleared after `getFileLocation`
Assertions.assertNull(CallerContext.CallerContextHolder.get());
}

private FilesetDTO mockFilesetDTO(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -291,6 +291,9 @@ public Response getFileLocation(
});
} catch (Exception e) {
return ExceptionHandlers.handleFilesetException(OperationType.GET, fileset, schema, e);
} finally {
// Clear the caller context
CallerContext.CallerContextHolder.remove();
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -554,6 +554,18 @@ public void testGetFileLocation() {
} finally {
CallerContext.CallerContextHolder.remove();
}

// test `getFileLocation` clear caller context finally
Map<String, String> testContextMap = Maps.newHashMap();
testContextMap.put("k", "v");
CallerContext context = CallerContext.builder().withContext(testContextMap).build();
CallerContext.CallerContextHolder.set(context);
FilesetOperations mockOperations = Mockito.mock(FilesetOperations.class);
Mockito.when(mockOperations.getFileLocation(any(), any(), any(), any(), any()))
.thenCallRealMethod();
mockOperations.getFileLocation(
"test_metalake", "test_catalog", "test_schema", "fileset4", "/test");
Assertions.assertNull(CallerContext.CallerContextHolder.get());
}

private void assertUpdateFileset(FilesetUpdatesRequest req, Fileset updatedFileset) {
Expand Down

0 comments on commit 8aee4a3

Please sign in to comment.