[EJBCLIENT-424] Refactor context data cleaning

wildfly · Feb 14, 2023 · 4531804 · 4531804
1 parent e62ea58
commit 4531804
Showing 1 changed file with 3 additions and 6 deletions.
diff --git a/src/main/java/org/jboss/ejb/protocol/remote/EJBClientChannel.java b/src/main/java/org/jboss/ejb/protocol/remote/EJBClientChannel.java
@@ -1320,13 +1320,10 @@ class MethodCallResultProducer implements EJBReceiverInvocationContext.ResultPro
              */
             private void cleanContextDataBeforeUnmarshalling(EJBClientInvocationContext clientInvocationContext) {
                 Map<String, Object> contextData = clientInvocationContext.getContextData();
-                @SuppressWarnings("unchecked")
-                Set<String> returnedContextDataKeys = (Set<String>) contextData.get(EJBClientInvocationContext.RETURNED_CONTEXT_DATA_KEY);
-                if(returnedContextDataKeys != null) {
-                    contextData.keySet().removeAll(returnedContextDataKeys);
-                }
+                contextData.keySet().removeIf(k -> (k != EJBClientInvocationContext.RETURNED_CONTEXT_DATA_KEY));
             }
-
+
+
             public Object apply(final Void ignored0, final Void ignored1) throws Exception {
                 final ResponseMessageInputStream response;
                 if(inputStream instanceof ResponseMessageInputStream) {