This is an automated email from the ASF dual-hosted git repository.
jshao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gravitino.git
The following commit(s) were added to refs/heads/main by this push:
new 8aee4a3bb [#4973] improvement(client-java/server): Clear
`CallerContext` in Java Client and Server when call `getFileLocation` finally
(#4974)
8aee4a3bb is described below
commit 8aee4a3bb0bb3e2d61afef16aa9176663bb7b47d
Author: xloya <[email protected]>
AuthorDate: Fri Sep 20 20:04:33 2024 +0800
[#4973] improvement(client-java/server): Clear `CallerContext` in Java
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.
---
.../apache/gravitino/client/FilesetCatalog.java | 33 +++++++++++++---------
.../gravitino/client/TestFilesetCatalog.java | 2 ++
.../server/web/rest/FilesetOperations.java | 3 ++
.../server/web/rest/TestFilesetOperations.java | 12 ++++++++
4 files changed, 36 insertions(+), 14 deletions(-)
diff --git
a/clients/client-java/src/main/java/org/apache/gravitino/client/FilesetCatalog.java
b/clients/client-java/src/main/java/org/apache/gravitino/client/FilesetCatalog.java
index e3b2aa9a6..a57482ff5 100644
---
a/clients/client-java/src/main/java/org/apache/gravitino/client/FilesetCatalog.java
+++
b/clients/client-java/src/main/java/org/apache/gravitino/client/FilesetCatalog.java
@@ -244,20 +244,25 @@ class FilesetCatalog extends BaseSchemaCatalog implements
org.apache.gravitino.f
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
diff --git
a/clients/client-java/src/test/java/org/apache/gravitino/client/TestFilesetCatalog.java
b/clients/client-java/src/test/java/org/apache/gravitino/client/TestFilesetCatalog.java
index 8a87fb72e..446af40ed 100644
---
a/clients/client-java/src/test/java/org/apache/gravitino/client/TestFilesetCatalog.java
+++
b/clients/client-java/src/test/java/org/apache/gravitino/client/TestFilesetCatalog.java
@@ -534,6 +534,8 @@ public class TestFilesetCatalog extends TestBase {
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(
diff --git
a/server/src/main/java/org/apache/gravitino/server/web/rest/FilesetOperations.java
b/server/src/main/java/org/apache/gravitino/server/web/rest/FilesetOperations.java
index 8dc22e022..b010f2609 100644
---
a/server/src/main/java/org/apache/gravitino/server/web/rest/FilesetOperations.java
+++
b/server/src/main/java/org/apache/gravitino/server/web/rest/FilesetOperations.java
@@ -291,6 +291,9 @@ public class FilesetOperations {
});
} catch (Exception e) {
return ExceptionHandlers.handleFilesetException(OperationType.GET,
fileset, schema, e);
+ } finally {
+ // Clear the caller context
+ CallerContext.CallerContextHolder.remove();
}
}
}
diff --git
a/server/src/test/java/org/apache/gravitino/server/web/rest/TestFilesetOperations.java
b/server/src/test/java/org/apache/gravitino/server/web/rest/TestFilesetOperations.java
index 4b5e1f864..62375dc4b 100644
---
a/server/src/test/java/org/apache/gravitino/server/web/rest/TestFilesetOperations.java
+++
b/server/src/test/java/org/apache/gravitino/server/web/rest/TestFilesetOperations.java
@@ -554,6 +554,18 @@ public class TestFilesetOperations extends JerseyTest {
} 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) {
