Fix exiting bal test when a program has listeners and running workers/async functions

compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/bir/codegen/JvmConstants.java

@@ -382,7 +382,7 @@ public class JvmConstants {
     public static final String LAMBDA_PREFIX = "$lambda$";
     public static final String POPULATE_METHOD_PREFIX = "$populate";
     public static final String ADD_METHOD = "add";
-    public static final String TEST_EXECUTION_STATE = "$testExecutionState";
+    public static final String TEST_EXECUTION_STATE = "__gH7W16nQmp0TestExecState__";
     public static final String GET_TEST_EXECUTION_STATE = "$getTestExecutionState";
 
     // scheduler related constants

compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/bir/codegen/JvmPackageGen.java

@@ -119,7 +119,6 @@
 import static org.wso2.ballerinalang.compiler.bir.codegen.JvmConstants.OBJECT;
 import static org.wso2.ballerinalang.compiler.bir.codegen.JvmConstants.SERVICE_EP_AVAILABLE;
 import static org.wso2.ballerinalang.compiler.bir.codegen.JvmConstants.TEST_EXECUTE_METHOD;
-import static org.wso2.ballerinalang.compiler.bir.codegen.JvmConstants.TEST_EXECUTION_STATE;
 import static org.wso2.ballerinalang.compiler.bir.codegen.JvmConstants.VALUE_CREATOR;
 import static org.wso2.ballerinalang.compiler.bir.codegen.JvmDesugarPhase.addDefaultableBooleanVarsToSignature;
 import static org.wso2.ballerinalang.compiler.bir.codegen.JvmDesugarPhase.rewriteRecordInits;
@@ -256,9 +255,6 @@ private static void generateStaticInitializer(ClassWriter cw, String className,
             setServiceEPAvailableField(cw, mv, serviceEPAvailable, className);
             setModuleStatusField(cw, mv, className);
             setCurrentModuleField(cw, mv, jvmConstantsGen, birPackage.packageID, className);
-            if (isTestablePackage) {
-                setInitialTestExecutionState(mv, className);
-            }
         }
         JvmCodeGenUtil.generateStrandMetadata(mv, className, birPackage.packageID, asyncDataCollector);
         mv.visitInsn(RETURN);
@@ -404,10 +400,6 @@ private void generateModuleClasses(BIRPackage module, Map<String, byte[]> jarEnt
                     }
                 }
 
-                if (isTestable) {
-                    generateTestExecutionStateField(cw);
-                }
-
                 MainMethodGen mainMethodGen = new MainMethodGen(symbolTable, jvmTypeGen, asyncDataCollector);
                 mainMethodGen.generateMainMethod(mainFunc, cw, module, moduleClass, serviceEPAvailable, isTestable);
                 initMethodGen.generateLambdaForModuleExecuteFunction(cw, moduleClass, jvmCastGen, mainFunc,
@@ -856,15 +848,4 @@ private boolean listenerDeclarationFound(BPackageSymbol packageSymbol) {
         }
         return false;
     }
-
-    private static void generateTestExecutionStateField(ClassWriter cw) {
-        FieldVisitor fv = cw.visitField(ACC_PUBLIC | ACC_STATIC, TEST_EXECUTION_STATE, "I",
-                null, null);
-        fv.visitEnd();
-    }
-
-    private static void setInitialTestExecutionState(MethodVisitor mv, String initClass) {
-        mv.visitInsn(ICONST_0);
-        mv.visitFieldInsn(PUTSTATIC, initClass, TEST_EXECUTION_STATE, "I");
-    }
 }

compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/bir/codegen/methodgen/InitMethodGen.java

@@ -71,7 +71,7 @@
 import static org.objectweb.asm.Opcodes.INVOKESPECIAL;
 import static org.objectweb.asm.Opcodes.INVOKESTATIC;
 import static org.objectweb.asm.Opcodes.INVOKEVIRTUAL;
-import static org.objectweb.asm.Opcodes.IRETURN;
+import static org.objectweb.asm.Opcodes.LRETURN;
 import static org.objectweb.asm.Opcodes.NEW;
 import static org.objectweb.asm.Opcodes.RETURN;
 import static org.wso2.ballerinalang.compiler.bir.codegen.JvmConstants.CREATE_TYPES_METHOD;
@@ -319,12 +319,10 @@
                 addCheckedInvocation(modExecFunc, pkg.packageID, MODULE_START_METHOD, retVarRef, boolRef);
 
         if (testExecuteFunc != null) {
-            lastBB = addCheckedInvocationWithArgs(modExecFunc, pkg.packageID, TEST_EXECUTE_METHOD, retVarRef, boolRef,
-                    functionArgs, Collections.emptyList());
-        }
-
-        if (!serviceEPAvailable && !JvmPackageGen.isLangModule(pkg.packageID)) {
-            lastBB = addCheckedInvocationForExitCall(modExecFunc, retVarRef, boolRef, typeOwnerClass);
+            lastBB = addTestExecuteInvocationWithGracefulExitCall(modExecFunc, pkg.packageID, retVarRef,
+                    functionArgs, Collections.emptyList(), typeOwnerClass);
+        } else if (!serviceEPAvailable && !JvmPackageGen.isLangModule(pkg.packageID)) {
+            lastBB = addInvocationForGracefulExitCall(modExecFunc, retVarRef, boolRef, typeOwnerClass);
         }
         lastBB.terminator = new BIRTerminator.Return(null);
         return modExecFunc;
@@ -446,21 +444,12 @@
         return new Name(varIdPrefix + nextVarId);
     }
 
-    private BIRNode.BIRBasicBlock addCheckedInvocationForExitCall(BIRNode.BIRFunction func, BIROperand retVar,
-                                                                  BIROperand boolRef, String typeOwnerClass) {
+    private BIRNode.BIRBasicBlock addInvocationForGracefulExitCall(BIRNode.BIRFunction func, BIROperand retVar,
+                                                                   BIROperand boolRef, String typeOwnerClass) {
         BIRNode.BIRBasicBlock lastBB = func.basicBlocks.get(func.basicBlocks.size() - 1);
         BIRNode.BIRBasicBlock nextBB = addAndGetNextBasicBlock(func);
         lastBB.terminator = getExitMethodCall(nextBB, typeOwnerClass);
-        BIRNonTerminator.TypeTest typeTest = new BIRNonTerminator.TypeTest(null, symbolTable.errorType, boolRef,
-                retVar);
-        nextBB.instructions.add(typeTest);
-        BIRNode.BIRBasicBlock trueBB = addAndGetNextBasicBlock(func);
-        BIRNode.BIRBasicBlock retBB = addAndGetNextBasicBlock(func);
-        retBB.terminator = new BIRTerminator.Return(null);
-        trueBB.terminator = new BIRTerminator.GOTO(null, retBB);
-        BIRNode.BIRBasicBlock falseBB = addAndGetNextBasicBlock(func);
-        nextBB.terminator = new BIRTerminator.Branch(null, boolRef, trueBB, falseBB);
-        return falseBB;
+        return nextBB;
     }
 
     private static JIMethodCall getExitMethodCall(BIRNode.BIRBasicBlock nextBB, String typeOwnerClass) {
@@ -476,6 +465,26 @@
         return jiMethodCall;
     }
 
+    private BIRNode.BIRBasicBlock addTestExecuteInvocationWithGracefulExitCall(BIRNode.BIRFunction func,
+                                                                               PackageID modId, BIROperand retVar,
+                                                                               List<BIROperand> args,
+                                                                               List<BIRNode.BIRAnnotationAttachment>
+                                                                                       calleeAnnotAttachments,
+                                                                               String typeOwnerClass) {
+        BIRNode.BIRBasicBlock lastBB = func.basicBlocks.get(func.basicBlocks.size() - 1);
+        BIRNode.BIRBasicBlock nextBB = addAndGetNextBasicBlock(func);
+        if (JvmCodeGenUtil.isBuiltInPackage(modId)) {
+            lastBB.terminator = new BIRTerminator.Call(null, InstructionKind.CALL, false, modId,
+                    new Name(TEST_EXECUTE_METHOD), args, null, nextBB, calleeAnnotAttachments, Collections.emptySet());
+            return nextBB;
+        }
+        lastBB.terminator = new BIRTerminator.Call(null, InstructionKind.CALL, false, modId,
+                new Name(TEST_EXECUTE_METHOD), args, retVar, nextBB, calleeAnnotAttachments, Collections.emptySet());
+        BIRNode.BIRBasicBlock finalBB = addAndGetNextBasicBlock(func);
+        nextBB.terminator = getExitMethodCall(finalBB, typeOwnerClass);
+        return finalBB;
+    }
+
     private BIRNode.BIRBasicBlock addCheckedInvocationWithArgs(BIRNode.BIRFunction func, PackageID modId,
                                                                String initFuncName,
                                                                BIROperand retVar, BIROperand boolRef,
@@ -540,11 +549,11 @@
     }
 
     public void generateGetTestExecutionState(ClassWriter cw, String className) {
-        MethodVisitor mv = cw.visitMethod(ACC_PUBLIC | ACC_STATIC, GET_TEST_EXECUTION_STATE, "()I",
+        MethodVisitor mv = cw.visitMethod(ACC_PUBLIC | ACC_STATIC, GET_TEST_EXECUTION_STATE, "()J",
                 null, null);
         mv.visitCode();
-        mv.visitFieldInsn(GETSTATIC, className, TEST_EXECUTION_STATE, "I");
-        mv.visitInsn(IRETURN);
+        mv.visitFieldInsn(GETSTATIC, className, TEST_EXECUTION_STATE, "J");
+        mv.visitInsn(LRETURN);
         JvmCodeGenUtil.visitMaxStackForMethod(mv, GET_TEST_EXECUTION_STATE, className);
         mv.visitEnd();
     }

compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/bir/codegen/methodgen/MainMethodGen.java

@@ -56,7 +56,6 @@
 import static org.objectweb.asm.Opcodes.INVOKEVIRTUAL;
 import static org.objectweb.asm.Opcodes.NEW;
 import static org.objectweb.asm.Opcodes.PUTFIELD;
-import static org.objectweb.asm.Opcodes.PUTSTATIC;
 import static org.objectweb.asm.Opcodes.RETURN;
 import static org.wso2.ballerinalang.compiler.bir.codegen.JvmConstants.CLI_SPEC;
 import static org.wso2.ballerinalang.compiler.bir.codegen.JvmConstants.COMPATIBILITY_CHECKER;
@@ -66,6 +65,7 @@
 import static org.wso2.ballerinalang.compiler.bir.codegen.JvmConstants.FUTURE_VALUE;
 import static org.wso2.ballerinalang.compiler.bir.codegen.JvmConstants.HANDLE_ALL_THROWABLE_METHOD;
 import static org.wso2.ballerinalang.compiler.bir.codegen.JvmConstants.HANDLE_RETURNED_ERROR_METHOD;
+import static org.wso2.ballerinalang.compiler.bir.codegen.JvmConstants.HANDLE_STOP_PANIC_METHOD;
 import static org.wso2.ballerinalang.compiler.bir.codegen.JvmConstants.HANDLE_THROWABLE_METHOD;
 import static org.wso2.ballerinalang.compiler.bir.codegen.JvmConstants.JAVA_RUNTIME;
 import static org.wso2.ballerinalang.compiler.bir.codegen.JvmConstants.JVM_INIT_METHOD;
@@ -90,7 +90,6 @@
 import static org.wso2.ballerinalang.compiler.bir.codegen.JvmConstants.STRING_VALUE;
 import static org.wso2.ballerinalang.compiler.bir.codegen.JvmConstants.TEST_ARGUMENTS;
 import static org.wso2.ballerinalang.compiler.bir.codegen.JvmConstants.TEST_CONFIG_ARGS;
-import static org.wso2.ballerinalang.compiler.bir.codegen.JvmConstants.TEST_EXECUTION_STATE;
 import static org.wso2.ballerinalang.compiler.bir.codegen.JvmConstants.THROWABLE;
 import static org.wso2.ballerinalang.compiler.bir.codegen.JvmConstants.TOML_DETAILS;
 import static org.wso2.ballerinalang.compiler.bir.codegen.JvmSignatures.ADD_SHUTDOWN_HOOK;
@@ -428,31 +427,30 @@ private void genSubmitToScheduler(String initClass, MethodVisitor mv, boolean is
         mv.visitFieldInsn(PUTFIELD, STRAND_CLASS, MethodGenUtils.FRAMES, STACK_FRAMES);
 
         startScheduler(indexMap.get(SCHEDULER_VAR), mv);
-        handleErrorFromFutureValue(mv, initClass, isTestFunction);
+        handleErrorFromFutureValue(mv, isTestFunction);
     }
 
     private void stopListeners(MethodVisitor mv, boolean isServiceEPAvailable) {
         mv.visitLdcInsn(isServiceEPAvailable);
         mv.visitMethodInsn(INVOKESTATIC , LAUNCH_UTILS, "stopListeners", "(Z)V", false);
     }
 
-    private void handleErrorFromFutureValue(MethodVisitor mv, String initClass, boolean isTestFunction) {
+    private void handleErrorFromFutureValue(MethodVisitor mv, boolean isTestFunction) {
         mv.visitVarInsn(ALOAD, indexMap.get(INIT_FUTURE_VAR));
         mv.visitInsn(DUP);
-        mv.visitFieldInsn(GETFIELD , FUTURE_VALUE , PANIC_FIELD,
-                          GET_THROWABLE);
+        mv.visitFieldInsn(GETFIELD , FUTURE_VALUE , PANIC_FIELD, GET_THROWABLE);
 
         // handle any runtime errors
         Label labelIf = new Label();
         mv.visitJumpInsn(IFNULL, labelIf);
-        if (!isTestFunction) {
+        if (isTestFunction) {
             mv.visitFieldInsn(GETFIELD , FUTURE_VALUE , PANIC_FIELD, GET_THROWABLE);
-            mv.visitMethodInsn(INVOKESTATIC, RUNTIME_UTILS, HANDLE_THROWABLE_METHOD, HANDLE_THROWABLE, false);
-            mv.visitInsn(RETURN);
+            mv.visitMethodInsn(INVOKESTATIC, RUNTIME_UTILS, HANDLE_STOP_PANIC_METHOD, HANDLE_THROWABLE, false);
         } else {
-            mv.visitInsn(ICONST_1);
-            mv.visitFieldInsn(PUTSTATIC, initClass, TEST_EXECUTION_STATE, "I");
+            mv.visitFieldInsn(GETFIELD, FUTURE_VALUE, PANIC_FIELD, GET_THROWABLE);
+            mv.visitMethodInsn(INVOKESTATIC, RUNTIME_UTILS, HANDLE_THROWABLE_METHOD, HANDLE_THROWABLE, false);
         }
+        mv.visitInsn(RETURN);
         mv.visitLabel(labelIf);
     }
 }

compiler/ballerina-lang/src/main/java/org/wso2/ballerinalang/compiler/bir/emit/InstructionEmitter.java

@@ -647,10 +647,30 @@
             case WAIT_ALL:
                 return emitWaitAll((BIRTerminator.WaitAll) term, tabs);
             default:
-                throw new IllegalStateException("Not a terminator instruction");
+                return emitBIRTerminator(term.getClass().getSimpleName(), term.lhsOp, term.getRhsOperands(), tabs,
+                        term.thenBB);
         }
     }
 
+    private static String emitBIRTerminator(String ins, BIROperand lhsOp, BIROperand[] rhsOperands, int tabs,
+                                            BIRNode.BIRBasicBlock thenBB) {
+        String str = "";
+        str += emitTabs(tabs);
+        str += emitVarRef(lhsOp);
+        str += emitSpaces(1);
+        str += "=";
+        str += emitSpaces(1);
+        str += ins;
+        str += emitSpaces(1);
+        str += emitVarRefs(rhsOperands);
+        str += emitSpaces(1);
+        str += "->";
+        str += emitSpaces(1);
+        str += emitBasicBlockRef(thenBB);
+        str += ";";
+        return str;
+    }
+
     private static String emitWait(BIRTerminator.Wait term, int tabs) {
 
         StringBuilder str = new StringBuilder();

misc/testerina/modules/testerina-compiler-plugin/src/main/java/org/ballerinalang/testerina/compiler/TestExecutionGenerationTask.java

@@ -81,6 +81,7 @@ private static TextDocument generateDocument(Module module) {
         List<ModuleMemberDeclarationNode> functionsList = new ArrayList<>();
         List<StatementNode> statements = new ArrayList<>();
 
+        TesterinaCompilerPluginUtils.addExitCodeGlobalVariable(functionsList);
         TesterinaCompilerPluginUtils.addSetTestOptionsCall(statements);
 
         // Initialize variables for test registrars

misc/testerina/modules/testerina-compiler-plugin/src/main/java/org/ballerinalang/testerina/compiler/TestExecutionModificationTask.java

@@ -63,6 +63,7 @@ private static TextDocument modifyDocument(Document document) {
         List<ModuleMemberDeclarationNode> functionsList = new ArrayList<>();
         List<StatementNode> statements = new ArrayList<>();
 
+        TesterinaCompilerPluginUtils.addExitCodeGlobalVariable(functionsList);
         TesterinaCompilerPluginUtils.addSetTestOptionsCall(statements);
 
         // TODO: replace visitor in modifier with a simple statement addition

misc/testerina/modules/testerina-compiler-plugin/src/main/java/org/ballerinalang/testerina/compiler/TesterinaCompilerPluginConstants.java

@@ -33,6 +33,7 @@ public class TesterinaCompilerPluginConstants {
     public static final String TEST_EXEC_FUNCTION = "__execute__";
     public static final String TEST_EXEC_FILENAME = "test_execute";
     public static final String TEST_REGISTRAR_EXEC_FUNCTION = "executeTestRegistrar";
+    public static final String TEST_EXECUTION_STATE = "__gH7W16nQmp0TestExecState__";
     public static final int REGISTERS_PER_FUNCTION = 150;
 
     public static final String TARGET_PATH_PARAMETER = "targetPath";
@@ -45,4 +46,6 @@ public class TesterinaCompilerPluginConstants {
     public static final String TESTS_PARAMETER = "tests";
     public static final String RERUN_FAILED_PARAMETER = "rerunFailed";
     public static final String LIST_GROUPS_PARAMETER = "listGroups";
+
+    private TesterinaCompilerPluginConstants() {}
 }