[EJBCLIENT-424] Refactor context data cleaning

wildfly · Mar 3, 2023 · e574cf2 · e574cf2
1 parent 958f4bb
commit e574cf2
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
@@ -1321,13 +1321,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.equals(EJBClientInvocationContext.RETURNED_CONTEXT_DATA_KEY)));
             }
-
+
+
             public Object apply(final Void ignored0, final Void ignored1) throws Exception {
                 final ResponseMessageInputStream response;
                 if(inputStream instanceof ResponseMessageInputStream) {