Merge "Separate stateful logic from JavaScriptIsolate" into androidx-main

NOKEYCHECK=True
GitOrigin-RevId: e56b0f3ddf1ad1ae259e5b0f3714ab8b71048390
diff --git a/main/java/androidx/javascriptengine/EnvironmentDeadState.java b/main/java/androidx/javascriptengine/EnvironmentDeadState.java
new file mode 100644
index 0000000..870d5ee
--- /dev/null
+++ b/main/java/androidx/javascriptengine/EnvironmentDeadState.java
@@ -0,0 +1,90 @@
+/*
+ * Copyright 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.javascriptengine;
+
+import androidx.annotation.NonNull;
+import androidx.concurrent.futures.CallbackToFutureAdapter;
+
+import com.google.common.util.concurrent.ListenableFuture;
+
+import java.util.concurrent.Executor;
+
+/**
+ * Covers the case where the environment is dead.
+ *
+ * This state covers cases where the developer explicitly closes the sandbox or sandbox/isolate
+ * being dead outside of the control of the developer.
+ */
+final class EnvironmentDeadState implements IsolateState {
+    private final JavaScriptException mException;
+
+    EnvironmentDeadState(JavaScriptException e) {
+        mException = e;
+    }
+
+    @NonNull
+    @Override
+    public ListenableFuture<String> evaluateJavaScriptAsync(@NonNull byte[] code) {
+        return CallbackToFutureAdapter.getFuture(completer -> {
+            final String futureDebugMessage = "evaluateJavascript Future";
+            completer.setException(mException);
+            return futureDebugMessage;
+        });
+    }
+
+    @NonNull
+    @Override
+    public ListenableFuture<String> evaluateJavaScriptAsync(@NonNull String code) {
+        return CallbackToFutureAdapter.getFuture(completer -> {
+            final String futureDebugMessage = "evaluateJavascript Future";
+            completer.setException(mException);
+            return futureDebugMessage;
+        });
+    }
+
+    @Override
+    public void setConsoleCallback(@NonNull Executor executor,
+            @NonNull JavaScriptConsoleCallback callback) {
+    }
+
+    @Override
+    public void setConsoleCallback(@NonNull JavaScriptConsoleCallback callback) {
+    }
+
+    @Override
+    public void clearConsoleCallback() {
+    }
+
+    @Override
+    public boolean provideNamedData(@NonNull String name, @NonNull byte[] inputBytes) {
+        return false;
+    }
+
+    @Override
+    public void close() {
+    }
+
+    @Override
+    public IsolateState setSandboxDead() {
+        return new EnvironmentDeadState(new SandboxDeadException());
+    }
+
+    @Override
+    public IsolateState setIsolateDead() {
+        return this;
+    }
+}
diff --git a/main/java/androidx/javascriptengine/IsolateClosedState.java b/main/java/androidx/javascriptengine/IsolateClosedState.java
new file mode 100644
index 0000000..0742719
--- /dev/null
+++ b/main/java/androidx/javascriptengine/IsolateClosedState.java
@@ -0,0 +1,84 @@
+/*
+ * Copyright 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.javascriptengine;
+
+import androidx.annotation.NonNull;
+
+import com.google.common.util.concurrent.ListenableFuture;
+
+import java.util.concurrent.Executor;
+
+/**
+ * Covers the case where the isolate is explicitly closed by the developer.
+ */
+final class IsolateClosedState implements IsolateState {
+    IsolateClosedState() {
+    }
+
+    @NonNull
+    @Override
+    public ListenableFuture<String> evaluateJavaScriptAsync(@NonNull byte[] code) {
+        throw new IllegalStateException("Calling evaluateJavaScriptAsync() after closing the"
+                + "Isolate");
+    }
+
+    @NonNull
+    @Override
+    public ListenableFuture<String> evaluateJavaScriptAsync(@NonNull String code) {
+        throw new IllegalStateException("Calling evaluateJavaScriptAsync() after closing the"
+                + "Isolate");
+    }
+
+    @Override
+    public void setConsoleCallback(@NonNull Executor executor,
+            @NonNull JavaScriptConsoleCallback callback) {
+        throw new IllegalStateException(
+                "Calling setConsoleCallback() after closing the Isolate");
+    }
+
+    @Override
+    public void setConsoleCallback(@NonNull JavaScriptConsoleCallback callback) {
+        throw new IllegalStateException(
+                "Calling setConsoleCallback() after closing the Isolate");
+    }
+
+    @Override
+    public void clearConsoleCallback() {
+        throw new IllegalStateException(
+                "Calling clearConsoleCallback() after closing the Isolate");
+    }
+
+    @Override
+    public boolean provideNamedData(@NonNull String name, @NonNull byte[] inputBytes) {
+        throw new IllegalStateException(
+                "Calling provideNamedData() after closing the Isolate");
+    }
+
+    @Override
+    public void close() {
+    }
+
+    @Override
+    public IsolateState setSandboxDead() {
+        return this;
+    }
+
+    @Override
+    public IsolateState setIsolateDead() {
+        return this;
+    }
+}
diff --git a/main/java/androidx/javascriptengine/IsolateState.java b/main/java/androidx/javascriptengine/IsolateState.java
new file mode 100644
index 0000000..a5fea09
--- /dev/null
+++ b/main/java/androidx/javascriptengine/IsolateState.java
@@ -0,0 +1,57 @@
+/*
+ * Copyright 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.javascriptengine;
+
+import androidx.annotation.NonNull;
+
+import com.google.common.util.concurrent.ListenableFuture;
+
+import java.util.concurrent.Executor;
+
+/**
+ * Interface for State design pattern.
+ *
+ * Isolates can be in different states due to events within/outside the control of the developer.
+ * This pattern allows us to extract out the state related behaviour without maintaining it all in
+ * the JavaScriptIsolate class which proved to be error-prone and hard to read.
+ *
+ * State specific behaviour are implemented in concrete classes that implements this interface.
+ *
+ * Refer: https://en.wikipedia.org/wiki/State_pattern
+ */
+interface IsolateState {
+    @NonNull
+    ListenableFuture<String> evaluateJavaScriptAsync(@NonNull byte[] code);
+
+    @NonNull
+    ListenableFuture<String> evaluateJavaScriptAsync(@NonNull String code);
+
+    void setConsoleCallback(@NonNull Executor executor,
+            @NonNull JavaScriptConsoleCallback callback);
+
+    void setConsoleCallback(@NonNull JavaScriptConsoleCallback callback);
+
+    void clearConsoleCallback();
+
+    boolean provideNamedData(@NonNull String name, @NonNull byte[] inputBytes);
+
+    void close();
+
+    IsolateState setIsolateDead();
+
+    IsolateState setSandboxDead();
+}
diff --git a/main/java/androidx/javascriptengine/IsolateUsableState.java b/main/java/androidx/javascriptengine/IsolateUsableState.java
new file mode 100644
index 0000000..6ae6286
--- /dev/null
+++ b/main/java/androidx/javascriptengine/IsolateUsableState.java
@@ -0,0 +1,410 @@
+/*
+ * Copyright 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.javascriptengine;
+
+import android.content.res.AssetFileDescriptor;
+import android.os.Binder;
+import android.os.DeadObjectException;
+import android.os.RemoteException;
+import android.util.Log;
+
+import androidx.annotation.NonNull;
+import androidx.concurrent.futures.CallbackToFutureAdapter;
+import androidx.javascriptengine.common.LengthLimitExceededException;
+import androidx.javascriptengine.common.Utils;
+
+import com.google.common.util.concurrent.ListenableFuture;
+
+import org.chromium.android_webview.js_sandbox.common.IJsSandboxConsoleCallback;
+import org.chromium.android_webview.js_sandbox.common.IJsSandboxIsolate;
+import org.chromium.android_webview.js_sandbox.common.IJsSandboxIsolateCallback;
+import org.chromium.android_webview.js_sandbox.common.IJsSandboxIsolateSyncCallback;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Objects;
+import java.util.Set;
+import java.util.concurrent.Executor;
+import java.util.concurrent.RejectedExecutionException;
+
+import javax.annotation.concurrent.GuardedBy;
+import javax.annotation.concurrent.NotThreadSafe;
+
+/**
+ * Covers the case where the isolate is functional.
+ */
+@NotThreadSafe
+final class IsolateUsableState implements IsolateState {
+    private static final String TAG = "IsolateUsableState";
+    final JavaScriptIsolate mJsIsolate;
+    private final Object mLock = new Object();
+    final int mMaxEvaluationReturnSizeBytes;
+
+    /**
+     * Interface to underlying service-backed implementation.
+     */
+    @NonNull
+    final IJsSandboxIsolate mJsIsolateStub;
+    @NonNull
+    @GuardedBy("mLock")
+    private Set<CallbackToFutureAdapter.Completer<String>> mPendingCompleterSet =
+            new HashSet<>();
+
+    private class IJsSandboxIsolateSyncCallbackStubWrapper extends
+            IJsSandboxIsolateSyncCallback.Stub {
+        @NonNull
+        private final CallbackToFutureAdapter.Completer<String> mCompleter;
+
+        IJsSandboxIsolateSyncCallbackStubWrapper(
+                @NonNull CallbackToFutureAdapter.Completer<String> completer) {
+            mCompleter = completer;
+        }
+
+        @Override
+        public void reportResultWithFd(AssetFileDescriptor afd) {
+            Objects.requireNonNull(afd);
+            mJsIsolate.mJsSandbox.mThreadPoolTaskExecutor.execute(
+                    () -> {
+                        String result;
+                        try {
+                            result = Utils.readToString(afd,
+                                    mMaxEvaluationReturnSizeBytes,
+                                    /*truncate=*/false);
+                        } catch (IOException | UnsupportedOperationException ex) {
+                            removePending(mCompleter);
+                            mCompleter.setException(
+                                    new JavaScriptException(
+                                            "Retrieving result failed: " + ex.getMessage()));
+                            return;
+                        } catch (LengthLimitExceededException ex) {
+                            removePending(mCompleter);
+                            if (ex.getMessage() != null) {
+                                mCompleter.setException(
+                                        new EvaluationResultSizeLimitExceededException(
+                                                ex.getMessage()));
+                            } else {
+                                mCompleter.setException(
+                                        new EvaluationResultSizeLimitExceededException());
+                            }
+                            return;
+                        }
+                        handleEvaluationResult(mCompleter, result);
+                    });
+        }
+
+        @Override
+        public void reportErrorWithFd(@ExecutionErrorTypes int type, AssetFileDescriptor afd) {
+            Objects.requireNonNull(afd);
+            mJsIsolate.mJsSandbox.mThreadPoolTaskExecutor.execute(
+                    () -> {
+                        String error;
+                        try {
+                            error = Utils.readToString(afd,
+                                    mMaxEvaluationReturnSizeBytes,
+                                    /*truncate=*/true);
+                        } catch (IOException | UnsupportedOperationException ex) {
+                            removePending(mCompleter);
+                            mCompleter.setException(
+                                    new JavaScriptException(
+                                            "Retrieving error failed: " + ex.getMessage()));
+                            return;
+                        } catch (LengthLimitExceededException ex) {
+                            throw new AssertionError("unreachable");
+                        }
+                        handleEvaluationError(mCompleter, type, error);
+                    });
+        }
+    }
+
+    private class IJsSandboxIsolateCallbackStubWrapper extends IJsSandboxIsolateCallback.Stub {
+        @NonNull
+        private final CallbackToFutureAdapter.Completer<String> mCompleter;
+
+        IJsSandboxIsolateCallbackStubWrapper(
+                @NonNull CallbackToFutureAdapter.Completer<String> completer) {
+            mCompleter = completer;
+        }
+
+        @Override
+        public void reportResult(String result) {
+            Objects.requireNonNull(result);
+            final long identityToken = Binder.clearCallingIdentity();
+            try {
+                handleEvaluationResult(mCompleter, result);
+            } finally {
+                Binder.restoreCallingIdentity(identityToken);
+            }
+        }
+
+        @Override
+        public void reportError(@ExecutionErrorTypes int type, String error) {
+            Objects.requireNonNull(error);
+            final long identityToken = Binder.clearCallingIdentity();
+            try {
+                handleEvaluationError(mCompleter, type, error);
+            } finally {
+                Binder.restoreCallingIdentity(identityToken);
+            }
+        }
+    }
+
+    static final class JsSandboxConsoleCallbackRelay
+            extends IJsSandboxConsoleCallback.Stub {
+        @NonNull
+        private final Executor mExecutor;
+        @NonNull
+        private final JavaScriptConsoleCallback mCallback;
+
+        JsSandboxConsoleCallbackRelay(@NonNull Executor executor,
+                @NonNull JavaScriptConsoleCallback callback) {
+            mExecutor = executor;
+            mCallback = callback;
+        }
+
+        @Override
+        public void consoleMessage(final int contextGroupId, final int level, final String message,
+                final String source, final int line, final int column, final String trace) {
+            final long identity = Binder.clearCallingIdentity();
+            try {
+                mExecutor.execute(() -> {
+                    if ((level & JavaScriptConsoleCallback.ConsoleMessage.LEVEL_ALL) == 0
+                            || ((level - 1) & level) != 0) {
+                        throw new IllegalArgumentException(
+                                "invalid console level " + level + " provided by isolate");
+                    }
+                    Objects.requireNonNull(message);
+                    Objects.requireNonNull(source);
+                    mCallback.onConsoleMessage(
+                            new JavaScriptConsoleCallback.ConsoleMessage(
+                                    level, message, source, line, column, trace));
+                });
+            } catch (RejectedExecutionException e) {
+                Log.e(TAG, "Console message dropped", e);
+            } finally {
+                Binder.restoreCallingIdentity(identity);
+            }
+        }
+
+        @Override
+        public void consoleClear(int contextGroupId) {
+            final long identity = Binder.clearCallingIdentity();
+            try {
+                mExecutor.execute(mCallback::onConsoleClear);
+            } catch (RejectedExecutionException e) {
+                Log.e(TAG, "Console clear dropped", e);
+            } finally {
+                Binder.restoreCallingIdentity(identity);
+            }
+        }
+
+    }
+
+    IsolateUsableState(JavaScriptIsolate isolate, @NonNull IJsSandboxIsolate jsIsolateStub,
+            int maxEvaluationResultSizeBytes) {
+        mJsIsolate = isolate;
+        mJsIsolateStub = jsIsolateStub;
+        mMaxEvaluationReturnSizeBytes = maxEvaluationResultSizeBytes;
+    }
+
+    @NonNull
+    @Override
+    public ListenableFuture<String> evaluateJavaScriptAsync(@NonNull byte[] code) {
+        return CallbackToFutureAdapter.getFuture(completer -> {
+            final String futureDebugMessage = "evaluateJavascript Future";
+            IJsSandboxIsolateSyncCallbackStubWrapper callbackStub =
+                    new IJsSandboxIsolateSyncCallbackStubWrapper(completer);
+            try {
+                // We pass the codeAfd to the separate sandbox process but we still need to
+                // close it on our end to avoid file descriptor leaks.
+                try (AssetFileDescriptor codeAfd = Utils.writeBytesIntoPipeAsync(code,
+                        mJsIsolate.mJsSandbox.mThreadPoolTaskExecutor)) {
+                    mJsIsolateStub.evaluateJavascriptWithFd(codeAfd,
+                            callbackStub);
+                }
+                addToPendingCompleterSet(completer);
+            } catch (DeadObjectException e) {
+                // The sandbox process has died.
+                mJsIsolate.maybeSetSandboxDead();
+                completer.setException(new SandboxDeadException());
+            } catch (RemoteException | IOException e) {
+                completer.setException(new RuntimeException(e));
+            }
+            // Debug string.
+            return futureDebugMessage;
+        });
+    }
+
+    @NonNull
+    @Override
+    public ListenableFuture<String> evaluateJavaScriptAsync(@NonNull String code) {
+        if (mJsIsolate.mJsSandbox.isFeatureSupported(
+                JavaScriptSandbox.JS_FEATURE_EVALUATE_WITHOUT_TRANSACTION_LIMIT)) {
+            // This process can be made more memory efficient by converting the
+            // String to UTF-8 encoded bytes and writing to the pipe in chunks.
+            byte[] inputBytes = code.getBytes(StandardCharsets.UTF_8);
+            return evaluateJavaScriptAsync(inputBytes);
+        }
+
+        return CallbackToFutureAdapter.getFuture(completer -> {
+            final String futureDebugMessage = "evaluateJavascript Future";
+            IJsSandboxIsolateCallbackStubWrapper callbackStub =
+                    new IJsSandboxIsolateCallbackStubWrapper(completer);
+            try {
+                mJsIsolateStub.evaluateJavascript(code, callbackStub);
+                addToPendingCompleterSet(completer);
+            } catch (DeadObjectException e) {
+                // The sandbox process has died.
+                mJsIsolate.maybeSetSandboxDead();
+                completer.setException(new SandboxDeadException());
+            } catch (RemoteException e) {
+                completer.setException(new RuntimeException(e));
+            }
+            // Debug string.
+            return futureDebugMessage;
+        });
+    }
+
+    @Override
+    public void setConsoleCallback(@NonNull Executor executor,
+            @NonNull JavaScriptConsoleCallback callback) {
+        try {
+            mJsIsolateStub.setConsoleCallback(
+                    new JsSandboxConsoleCallbackRelay(executor, callback));
+        } catch (DeadObjectException e) {
+            mJsIsolate.maybeSetSandboxDead();
+        } catch (RemoteException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public void setConsoleCallback(@NonNull JavaScriptConsoleCallback callback) {
+        setConsoleCallback(mJsIsolate.mJsSandbox.getMainExecutor(), callback);
+    }
+
+    @Override
+    public void clearConsoleCallback() {
+        try {
+            mJsIsolateStub.setConsoleCallback(null);
+        } catch (DeadObjectException e) {
+            mJsIsolate.maybeSetSandboxDead();
+        } catch (RemoteException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public boolean provideNamedData(@NonNull String name, @NonNull byte[] inputBytes) {
+        try {
+            // We pass the codeAfd to the separate sandbox process but we still need to close
+            // it on our end to avoid file descriptor leaks.
+            try (AssetFileDescriptor codeAfd = Utils.writeBytesIntoPipeAsync(inputBytes,
+                    mJsIsolate.mJsSandbox.mThreadPoolTaskExecutor)) {
+                return mJsIsolateStub.provideNamedData(name, codeAfd);
+            }
+        } catch (DeadObjectException e) {
+            mJsIsolate.maybeSetSandboxDead();
+        } catch (RemoteException e) {
+            Log.e(TAG, "RemoteException was thrown during provideNamedData()", e);
+        } catch (IOException e) {
+            Log.e(TAG, "IOException was thrown during provideNamedData", e);
+        }
+        return false;
+    }
+
+    @Override
+    public void close() {
+        try {
+            mJsIsolateStub.close();
+        } catch (DeadObjectException e) {
+            Log.e(TAG, "DeadObjectException was thrown during close()", e);
+        } catch (RemoteException e) {
+            Log.e(TAG, "RemoteException was thrown during close()", e);
+        }
+        cancelAllPendingEvaluations(new IsolateTerminatedException());
+    }
+
+    @Override
+    public IsolateState setIsolateDead() {
+        IsolateTerminatedException exception = new IsolateTerminatedException();
+        cancelAllPendingEvaluations(exception);
+        return new EnvironmentDeadState(exception);
+    }
+
+    @Override
+    public IsolateState setSandboxDead() {
+        SandboxDeadException exception = new SandboxDeadException();
+        cancelAllPendingEvaluations(exception);
+        return new EnvironmentDeadState(exception);
+    }
+
+    void handleEvaluationError(@NonNull CallbackToFutureAdapter.Completer<String> completer,
+            int type, @NonNull String error) {
+        removePending(completer);
+        boolean crashing = false;
+        switch (type) {
+            case IJsSandboxIsolateSyncCallback.JS_EVALUATION_ERROR:
+                completer.setException(new EvaluationFailedException(error));
+                break;
+            case IJsSandboxIsolateSyncCallback.MEMORY_LIMIT_EXCEEDED:
+                completer.setException(new MemoryLimitExceededException(error));
+                crashing = true;
+                break;
+            default:
+                completer.setException(new JavaScriptException(
+                        "Crashing due to unknown JavaScriptException: " + error));
+                // Assume the worst
+                crashing = true;
+        }
+        if (crashing) {
+            mJsIsolate.maybeSetIsolateDead();
+        }
+    }
+
+    void handleEvaluationResult(@NonNull CallbackToFutureAdapter.Completer<String> completer,
+            @NonNull String result) {
+        removePending(completer);
+        completer.set(result);
+    }
+
+    boolean removePending(@NonNull CallbackToFutureAdapter.Completer<String> completer) {
+        synchronized (mLock) {
+            return mPendingCompleterSet.remove(completer);
+        }
+    }
+
+    void addToPendingCompleterSet(@NonNull CallbackToFutureAdapter.Completer<String> completer) {
+        synchronized (mLock) {
+            mPendingCompleterSet.add(completer);
+        }
+    }
+
+    // Cancel all pending and future evaluations with the given exception.
+    // Only the first call to this method has any effect.
+    void cancelAllPendingEvaluations(@NonNull Exception e) {
+        Set<CallbackToFutureAdapter.Completer<String>> completers;
+        synchronized (mLock) {
+            completers = mPendingCompleterSet;
+            mPendingCompleterSet = Collections.emptySet();
+        }
+        for (CallbackToFutureAdapter.Completer<String> ele : completers) {
+            ele.setException(e);
+        }
+    }
+}
diff --git a/main/java/androidx/javascriptengine/JavaScriptIsolate.java b/main/java/androidx/javascriptengine/JavaScriptIsolate.java
index 9a8f735..c07746b 100644
--- a/main/java/androidx/javascriptengine/JavaScriptIsolate.java
+++ b/main/java/androidx/javascriptengine/JavaScriptIsolate.java
@@ -16,32 +16,15 @@
 
 package androidx.javascriptengine;
 
-import android.content.res.AssetFileDescriptor;
-import android.os.Binder;
-import android.os.RemoteException;
-import android.util.Log;
-
 import androidx.annotation.NonNull;
-import androidx.annotation.Nullable;
 import androidx.annotation.RequiresFeature;
-import androidx.concurrent.futures.CallbackToFutureAdapter;
-import androidx.javascriptengine.common.LengthLimitExceededException;
-import androidx.javascriptengine.common.Utils;
 
 import com.google.common.util.concurrent.ListenableFuture;
 
-import org.chromium.android_webview.js_sandbox.common.IJsSandboxConsoleCallback;
 import org.chromium.android_webview.js_sandbox.common.IJsSandboxIsolate;
-import org.chromium.android_webview.js_sandbox.common.IJsSandboxIsolateCallback;
-import org.chromium.android_webview.js_sandbox.common.IJsSandboxIsolateSyncCallback;
 
-import java.io.IOException;
-import java.nio.charset.StandardCharsets;
-import java.util.HashSet;
 import java.util.Objects;
 import java.util.concurrent.Executor;
-import java.util.concurrent.RejectedExecutionException;
-import java.util.concurrent.atomic.AtomicBoolean;
 
 import javax.annotation.concurrent.GuardedBy;
 
@@ -59,194 +42,55 @@
  * isolation multiple {@link JavaScriptSandbox} processes should be used, but it is not supported
  * at the moment. Please find the feature request <a href="https://crbug.com/1349860">here</a>.
  * <p>
- * Each isolate object must only be used from one thread.
+ * This class is thread-safe.
  */
 public final class JavaScriptIsolate implements AutoCloseable {
     private static final String TAG = "JavaScriptIsolate";
-    private final Object mSetLock = new Object();
-    /**
-     * Interface to underlying service-backed implementation.
-     * <p>
-     * mJsIsolateStub should only be null when the Isolate has been explicitly closed - not when the
-     * isolate has crashed or simply had its pending and future evaluations cancelled.
-     */
-    @Nullable
-    private IJsSandboxIsolate mJsIsolateStub;
+    private final Object mLock = new Object();
     private final CloseGuardHelper mGuard = CloseGuardHelper.create();
+
     final JavaScriptSandbox mJsSandbox;
 
-    @Nullable
-    @GuardedBy("mSetLock")
-    private HashSet<CallbackToFutureAdapter.Completer<String>> mPendingCompleterSet =
-            new HashSet<>();
-    /**
-     * If mSandboxClosed is true, new evaluations will throw this exception asynchronously.
-     * <p>
-     * Note that if the isolate is closed, IllegalStateException is thrown synchronously instead.
-     */
-    @Nullable
-    private Exception mExceptionForNewEvaluations;
-    private final AtomicBoolean mSandboxClosed = new AtomicBoolean(false);
-    final IsolateStartupParameters mStartupParameters;
-
-    private class IJsSandboxIsolateSyncCallbackStubWrapper extends
-            IJsSandboxIsolateSyncCallback.Stub {
-        @NonNull
-        private final CallbackToFutureAdapter.Completer<String> mCompleter;
-
-        IJsSandboxIsolateSyncCallbackStubWrapper(
-                @NonNull CallbackToFutureAdapter.Completer<String> completer) {
-            mCompleter = completer;
-        }
-
-        @Override
-        public void reportResultWithFd(AssetFileDescriptor afd) {
-            Objects.requireNonNull(afd);
-            mJsSandbox.mThreadPoolTaskExecutor.execute(
-                    () -> {
-                        String result;
-                        try {
-                            result = Utils.readToString(afd,
-                                    mStartupParameters.getMaxEvaluationReturnSizeBytes(),
-                                    /*truncate=*/false);
-                        } catch (IOException | UnsupportedOperationException ex) {
-                            mCompleter.setException(
-                                    new JavaScriptException(
-                                            "Retrieving result failed: " + ex.getMessage()));
-                            removePending(mCompleter);
-                            return;
-                        } catch (LengthLimitExceededException ex) {
-                            if (ex.getMessage() != null) {
-                                mCompleter.setException(
-                                        new EvaluationResultSizeLimitExceededException(
-                                                ex.getMessage()));
-                            } else {
-                                mCompleter.setException(
-                                        new EvaluationResultSizeLimitExceededException());
-                            }
-                            removePending(mCompleter);
-                            return;
-                        }
-                        handleEvaluationResult(mCompleter, result);
-                    });
-        }
-
-        @Override
-        public void reportErrorWithFd(@ExecutionErrorTypes int type, AssetFileDescriptor afd) {
-            Objects.requireNonNull(afd);
-            mJsSandbox.mThreadPoolTaskExecutor.execute(
-                    () -> {
-                        String error;
-                        try {
-                            error = Utils.readToString(afd,
-                                    mStartupParameters.getMaxEvaluationReturnSizeBytes(),
-                                    /*truncate=*/true);
-                        } catch (IOException | UnsupportedOperationException ex) {
-                            mCompleter.setException(
-                                    new JavaScriptException(
-                                            "Retrieving error failed: " + ex.getMessage()));
-                            removePending(mCompleter);
-                            return;
-                        } catch (LengthLimitExceededException ex) {
-                            throw new AssertionError("unreachable");
-                        }
-                        handleEvaluationError(mCompleter, type, error);
-                    });
-        }
-    }
-
-    private class IJsSandboxIsolateCallbackStubWrapper extends IJsSandboxIsolateCallback.Stub {
-        @NonNull
-        private final CallbackToFutureAdapter.Completer<String> mCompleter;
-
-        IJsSandboxIsolateCallbackStubWrapper(
-                @NonNull CallbackToFutureAdapter.Completer<String> completer) {
-            mCompleter = completer;
-        }
-
-        @Override
-        public void reportResult(String result) {
-            Objects.requireNonNull(result);
-            final long identityToken = Binder.clearCallingIdentity();
-            try {
-                handleEvaluationResult(mCompleter, result);
-            } finally {
-                Binder.restoreCallingIdentity(identityToken);
-            }
-        }
-
-        @Override
-        public void reportError(@ExecutionErrorTypes int type, String error) {
-            Objects.requireNonNull(error);
-            final long identityToken = Binder.clearCallingIdentity();
-            try {
-                handleEvaluationError(mCompleter, type, error);
-            } finally {
-                Binder.restoreCallingIdentity(identityToken);
-            }
-        }
-    }
-
-    private static final class JsSandboxConsoleCallbackRelay
-            extends IJsSandboxConsoleCallback.Stub {
-        @NonNull
-        private final Executor mExecutor;
-        @NonNull
-        private final JavaScriptConsoleCallback mCallback;
-
-        JsSandboxConsoleCallbackRelay(@NonNull Executor executor,
-                @NonNull JavaScriptConsoleCallback callback) {
-            mExecutor = executor;
-            mCallback = callback;
-        }
-
-        @Override
-        public void consoleMessage(final int contextGroupId, final int level, final String message,
-                final String source, final int line, final int column, final String trace) {
-            final long identity = Binder.clearCallingIdentity();
-            try {
-                mExecutor.execute(() -> {
-                    if ((level & JavaScriptConsoleCallback.ConsoleMessage.LEVEL_ALL) == 0
-                            || ((level - 1) & level) != 0) {
-                        throw new IllegalArgumentException(
-                                "invalid console level " + level + " provided by isolate");
-                    }
-                    Objects.requireNonNull(message);
-                    Objects.requireNonNull(source);
-                    mCallback.onConsoleMessage(
-                            new JavaScriptConsoleCallback.ConsoleMessage(
-                                    level, message, source, line, column, trace));
-                });
-            } catch (RejectedExecutionException e) {
-                Log.e(TAG, "Console message dropped", e);
-            } finally {
-                Binder.restoreCallingIdentity(identity);
-            }
-        }
-
-        @Override
-        public void consoleClear(int contextGroupId) {
-            final long identity = Binder.clearCallingIdentity();
-            try {
-                mExecutor.execute(mCallback::onConsoleClear);
-            } catch (RejectedExecutionException e) {
-                Log.e(TAG, "Console clear dropped", e);
-            } finally {
-                Binder.restoreCallingIdentity(identity);
-            }
-        }
-    }
+    @GuardedBy("mLock")
+    @NonNull
+    private IsolateState mIsolateState;
 
     JavaScriptIsolate(@NonNull IJsSandboxIsolate jsIsolateStub, @NonNull JavaScriptSandbox sandbox,
             @NonNull IsolateStartupParameters settings) {
+        synchronized (mLock) {
+            mIsolateState = new IsolateUsableState(this, jsIsolateStub,
+                    settings.getMaxEvaluationReturnSizeBytes());
+        }
         mJsSandbox = sandbox;
-        mJsIsolateStub = jsIsolateStub;
-        mStartupParameters = settings;
         mGuard.open("close");
         // This should be at the end of the constructor.
     }
 
     /**
+     * Changes the state to denote that the isolate is dead.
+     *
+     * {@link IsolateClosedState} takes precedence so it will not change state if the current state
+     * is {@link IsolateClosedState}
+     */
+    void maybeSetIsolateDead() {
+        synchronized (mLock) {
+            mIsolateState = mIsolateState.setIsolateDead();
+        }
+    }
+
+    /**
+     * Changes the state to denote that the sandbox is dead.
+     *
+     * {@link IsolateClosedState} takes precedence so it will not change state if the current state
+     * is {@link IsolateClosedState}
+     */
+    void maybeSetSandboxDead() {
+        synchronized (mLock) {
+            mIsolateState = mIsolateState.setSandboxDead();
+        }
+    }
+
+    /**
      * Evaluates the given JavaScript code and returns the result.
      * <p>
      * There are 3 possible behaviors based on the output of the expression:
@@ -281,50 +125,19 @@
      * <p>
      * Do not use this method to transfer raw binary data. Scripts or results containing unpaired
      * surrogate code units are not supported.
+     *
      * @param code JavaScript code to evaluate. The script should return a JavaScript String or,
      *             alternatively, a Promise that will resolve to a String if
      *             {@link JavaScriptSandbox#JS_FEATURE_PROMISE_RETURN} is supported.
      * @return Future that evaluates to the result String of the evaluation or exceptions (see
      * {@link JavaScriptException} and subclasses) if there is an error.
      */
-    @SuppressWarnings("NullAway")
     @NonNull
     public ListenableFuture<String> evaluateJavaScriptAsync(@NonNull String code) {
         Objects.requireNonNull(code);
-        if (!mSandboxClosed.get() && mJsSandbox.isFeatureSupported(
-                JavaScriptSandbox.JS_FEATURE_EVALUATE_WITHOUT_TRANSACTION_LIMIT)) {
-            // This process can be made more memory efficient by converting the String to
-            // UTF-8 encoded bytes and writing to the pipe in chunks.
-            byte[] inputBytes = code.getBytes(StandardCharsets.UTF_8);
-            return evaluateJavaScriptAsync(inputBytes);
+        synchronized (mLock) {
+            return mIsolateState.evaluateJavaScriptAsync(code);
         }
-        if (mJsIsolateStub == null) {
-            throw new IllegalStateException(
-                    "Calling evaluateJavaScriptAsync() after closing the Isolate");
-        }
-        return CallbackToFutureAdapter.getFuture(completer -> {
-            final String futureDebugMessage = "evaluateJavascript Future";
-            IJsSandboxIsolateCallbackStubWrapper callbackStub;
-            synchronized (mSetLock) {
-                if (mPendingCompleterSet == null) {
-                    assert mExceptionForNewEvaluations != null;
-                    completer.setException(mExceptionForNewEvaluations);
-                    return futureDebugMessage;
-                }
-                mPendingCompleterSet.add(completer);
-            }
-            callbackStub = new IJsSandboxIsolateCallbackStubWrapper(completer);
-            try {
-                mJsIsolateStub.evaluateJavascript(code, callbackStub);
-            } catch (RemoteException e) {
-                completer.setException(new RuntimeException(e));
-                synchronized (mSetLock) {
-                    mPendingCompleterSet.remove(completer);
-                }
-            }
-            // Debug string.
-            return futureDebugMessage;
-        });
     }
 
     /**
@@ -338,6 +151,7 @@
      * This overload is provided for clients to pass in a UTF-8 encoded {@code byte[]} directly
      * instead of having to convert it into a {@code String} to use
      * {@link #evaluateJavaScriptAsync(String)}.
+     *
      * @param code UTF-8 encoded JavaScript code to evaluate. The script should return a
      *             JavaScript String or, alternatively, a Promise that will resolve to a String if
      *             {@link JavaScriptSandbox#JS_FEATURE_PROMISE_RETURN} is supported. The byte
@@ -345,46 +159,14 @@
      * @return Future that evaluates to the result String of the evaluation or exceptions (see
      * {@link JavaScriptException} and subclasses) if there is an error
      */
-    @SuppressWarnings("NullAway")
     @NonNull
     @RequiresFeature(name = JavaScriptSandbox.JS_FEATURE_EVALUATE_WITHOUT_TRANSACTION_LIMIT,
             enforcement = "androidx.javascriptengine.JavaScriptSandbox#isFeatureSupported")
     public ListenableFuture<String> evaluateJavaScriptAsync(@NonNull byte[] code) {
         Objects.requireNonNull(code);
-        if (mJsIsolateStub == null) {
-            throw new IllegalStateException(
-                    "Calling evaluateJavaScriptAsync() after closing the Isolate");
+        synchronized (mLock) {
+            return mIsolateState.evaluateJavaScriptAsync(code);
         }
-        return CallbackToFutureAdapter.getFuture(completer -> {
-            final String futureDebugMessage = "evaluateJavascript Future";
-            IJsSandboxIsolateSyncCallbackStubWrapper callbackStub;
-            synchronized (mSetLock) {
-                if (mPendingCompleterSet == null) {
-                    assert mExceptionForNewEvaluations != null;
-                    completer.setException(mExceptionForNewEvaluations);
-                    return futureDebugMessage;
-                }
-                mPendingCompleterSet.add(completer);
-            }
-            callbackStub = new IJsSandboxIsolateSyncCallbackStubWrapper(completer);
-            try {
-                // We pass the codeAfd to the separate sandbox process but we still need to
-                // close it on our end to avoid file descriptor leaks.
-                try (AssetFileDescriptor codeAfd = Utils.writeBytesIntoPipeAsync(code,
-                        mJsSandbox.mThreadPoolTaskExecutor)) {
-                    mJsIsolateStub.evaluateJavascriptWithFd(codeAfd, callbackStub);
-                }
-            } catch (RemoteException | IOException e) {
-                completer.setException(new RuntimeException(e));
-                synchronized (mSetLock) {
-                    if (mPendingCompleterSet != null) {
-                        mPendingCompleterSet.remove(completer);
-                    }
-                }
-            }
-            // Debug string.
-            return futureDebugMessage;
-        });
     }
 
     /**
@@ -401,20 +183,12 @@
      */
     @Override
     public void close() {
-        // IllegalStateException will be thrown synchronously instead for new evaluations.
-        mExceptionForNewEvaluations = null;
-        if (mJsIsolateStub == null) {
-            return;
+        synchronized (mLock) {
+            mIsolateState.close();
+            mIsolateState = new IsolateClosedState();
+            mJsSandbox.removeFromIsolateSet(this);
+            mGuard.close();
         }
-        try {
-            cancelAllPendingEvaluations(new IsolateTerminatedException());
-            mJsIsolateStub.close();
-        } catch (RemoteException e) {
-            Log.e(TAG, "RemoteException was thrown during close()", e);
-        }
-        mJsIsolateStub = null;
-        mJsSandbox.removeFromIsolateSet(this);
-        mGuard.close();
     }
 
     /**
@@ -470,83 +244,9 @@
     public boolean provideNamedData(@NonNull String name, @NonNull byte[] inputBytes) {
         Objects.requireNonNull(name);
         Objects.requireNonNull(inputBytes);
-        if (mJsIsolateStub == null) {
-            throw new IllegalStateException("Calling provideNamedData() after closing the Isolate");
+        synchronized (mLock) {
+            return mIsolateState.provideNamedData(name, inputBytes);
         }
-        try {
-            // We pass the codeAfd to the separate sandbox process but we still need to close
-            // it on our end to avoid file descriptor leaks.
-            try (AssetFileDescriptor codeAfd = Utils.writeBytesIntoPipeAsync(inputBytes,
-                    mJsSandbox.mThreadPoolTaskExecutor)) {
-                return mJsIsolateStub.provideNamedData(name, codeAfd);
-            }
-        } catch (RemoteException e) {
-            Log.e(TAG, "RemoteException was thrown during provideNamedData()", e);
-        } catch (IOException e) {
-            Log.e(TAG, "IOException was thrown during provideNamedData", e);
-        }
-        return false;
-    }
-
-    void handleEvaluationError(@NonNull CallbackToFutureAdapter.Completer<String> completer,
-            int type, @NonNull String error) {
-        boolean crashing = false;
-        switch (type) {
-            case IJsSandboxIsolateSyncCallback.JS_EVALUATION_ERROR:
-                completer.setException(new EvaluationFailedException(error));
-                break;
-            case IJsSandboxIsolateSyncCallback.MEMORY_LIMIT_EXCEEDED:
-                completer.setException(new MemoryLimitExceededException(error));
-                crashing = true;
-                break;
-            default:
-                completer.setException(new JavaScriptException(
-                        "Crashing due to unknown JavaScriptException: " + error));
-                // Assume the worst
-                crashing = true;
-        }
-        removePending(completer);
-        if (crashing) {
-            handleCrash();
-        }
-    }
-
-    void handleEvaluationResult(@NonNull CallbackToFutureAdapter.Completer<String> completer,
-            @NonNull String result) {
-        completer.set(result);
-        removePending(completer);
-    }
-
-    void notifySandboxClosed() {
-        mSandboxClosed.set(true);
-        cancelAllPendingEvaluations(new SandboxDeadException());
-    }
-
-    // Cancel all pending and future evaluations with the given exception.
-    // Only the first call to this method has any effect.
-    void cancelAllPendingEvaluations(@NonNull Exception e) {
-        final HashSet<CallbackToFutureAdapter.Completer<String>> pendingSet;
-        synchronized (mSetLock) {
-            if (mPendingCompleterSet == null) return;
-            pendingSet = mPendingCompleterSet;
-            mPendingCompleterSet = null;
-            mExceptionForNewEvaluations = e;
-        }
-        for (CallbackToFutureAdapter.Completer<String> ele : pendingSet) {
-            ele.setException(e);
-        }
-    }
-
-    void removePending(@NonNull CallbackToFutureAdapter.Completer<String> completer) {
-        synchronized (mSetLock) {
-            if (mPendingCompleterSet != null) {
-                mPendingCompleterSet.remove(completer);
-            }
-        }
-    }
-
-    void handleCrash() {
-        cancelAllPendingEvaluations(new IsolateTerminatedException());
     }
 
     @Override
@@ -554,9 +254,7 @@
     protected void finalize() throws Throwable {
         try {
             mGuard.warnIfOpen();
-            if (mJsIsolateStub != null) {
-                close();
-            }
+            close();
         } finally {
             super.finalize();
         }
@@ -592,15 +290,8 @@
             @NonNull JavaScriptConsoleCallback callback) {
         Objects.requireNonNull(executor);
         Objects.requireNonNull(callback);
-        if (mJsIsolateStub == null) {
-            throw new IllegalStateException(
-                    "Calling setConsoleCallback() after closing the Isolate");
-        }
-        try {
-            mJsIsolateStub.setConsoleCallback(
-                    new JsSandboxConsoleCallbackRelay(executor, callback));
-        } catch (RemoteException e) {
-            throw new RuntimeException(e);
+        synchronized (mLock) {
+            mIsolateState.setConsoleCallback(executor, callback);
         }
     }
 
@@ -609,13 +300,17 @@
      * <p>
      * This is the same as calling {@link #setConsoleCallback(Executor, JavaScriptConsoleCallback)}
      * using the main executor of the context used to create the {@link JavaScriptSandbox} object.
+     *
      * @param callback Callback implementing console logging behaviour.
      */
     @RequiresFeature(name = JavaScriptSandbox.JS_FEATURE_CONSOLE_MESSAGING,
             enforcement = "androidx.javascriptengine.JavaScriptSandbox#isFeatureSupported")
     public void setConsoleCallback(@NonNull JavaScriptConsoleCallback callback) {
         Objects.requireNonNull(callback);
-        setConsoleCallback(mJsSandbox.getMainExecutor(), callback);
+        synchronized (mLock) {
+            mIsolateState.setConsoleCallback(mJsSandbox.getMainExecutor(),
+                    callback);
+        }
     }
 
     /**
@@ -626,14 +321,8 @@
     @RequiresFeature(name = JavaScriptSandbox.JS_FEATURE_CONSOLE_MESSAGING,
             enforcement = "androidx.javascriptengine.JavaScriptSandbox#isFeatureSupported")
     public void clearConsoleCallback() {
-        if (mJsIsolateStub == null) {
-            throw new IllegalStateException(
-                    "Calling clearConsoleCallback() after closing the Isolate");
-        }
-        try {
-            mJsIsolateStub.setConsoleCallback(null);
-        } catch (RemoteException e) {
-            throw new RuntimeException(e);
+        synchronized (mLock) {
+            mIsolateState.clearConsoleCallback();
         }
     }
 }
diff --git a/main/java/androidx/javascriptengine/JavaScriptSandbox.java b/main/java/androidx/javascriptengine/JavaScriptSandbox.java
index b593986..834502a 100644
--- a/main/java/androidx/javascriptengine/JavaScriptSandbox.java
+++ b/main/java/androidx/javascriptengine/JavaScriptSandbox.java
@@ -531,7 +531,7 @@
     @GuardedBy("mLock")
     private void notifyIsolatesAboutClosureLocked() {
         for (JavaScriptIsolate ele : mActiveIsolateSet) {
-            ele.notifySandboxClosed();
+            ele.maybeSetSandboxDead();
         }
         mActiveIsolateSet.clear();
     }