Improve callMethod logic

ballerina-platform · Nov 28, 2024 · c76db7c · c76db7c
1 parent b9c5d03
commit c76db7c
Showing 1 changed file with 25 additions and 20 deletions.
diff --git a/native/src/main/java/io/ballerina/lib/data/xmldata/io/BallerinaByteBlockInputStream.java b/native/src/main/java/io/ballerina/lib/data/xmldata/io/BallerinaByteBlockInputStream.java
@@ -82,35 +82,40 @@ public int read() {
     @Override
     public void close() throws IOException {
         super.close();
-        if (closeMethod != null) {
-            env.getRuntime().callMethod(iterator, closeMethod.getName(), null);
-        }
+        env.yieldAndRun(() -> {
+            if (closeMethod != null) {
+                    env.getRuntime().callMethod(iterator, closeMethod.getName(), null);
+            }
+            return null;
+        });
     }
 
     private boolean hasBytesInCurrentChunk() {
         return currentChunk.length != 0 && nextChunkIndex < currentChunk.length;
     }
 
     private boolean readNextChunk() throws InterruptedException {
-        try {
-            Object result = env.getRuntime().callMethod(iterator, nextMethodName, null);
-            if (result == null) {
+        return env.yieldAndRun(() -> {
+            try {
+                Object result = env.getRuntime().callMethod(iterator, nextMethodName, null);
+                if (result == null) {
+                    done.set(true);
+                    return !done.get();
+                }
+                if (result instanceof BMap<?, ?>) {
+                    BMap<BString, Object> valueRecord = (BMap<BString, Object>) result;
+                    final BString value = Arrays.stream(valueRecord.getKeys()).findFirst().get();
+                    final BArray arrayValue = valueRecord.getArrayValue(value);
+                    currentChunk = arrayValue.getByteArray();
+                } else {
+                    done.set(true);
+                }
+            } catch (BError bError) {
                 done.set(true);
-                return true;
+                currentChunk = new byte[0];
             }
-            if (result instanceof BMap<?, ?>) {
-                BMap<BString, Object> valueRecord = (BMap<BString, Object>) result;
-                final BString value = Arrays.stream(valueRecord.getKeys()).findFirst().get();
-                final BArray arrayValue = valueRecord.getArrayValue(value);
-                currentChunk = arrayValue.getByteArray();
-            } else {
-                done.set(true);
-            }
-        } catch (BError bError) {
-            done.set(true);
-            currentChunk = new byte[0];
-        }
-        return !done.get();
+            return !done.get();
+        });
     }
 
     public BError getError() {