2 files changed, 101 insertions, 4 deletions

diff --git a/core/java/android/view/translation/TranslationManager.java b/core/java/android/view/translation/TranslationManager.java
index e75577487d1c..20d817db02b3 100644
--- a/core/java/android/view/translation/TranslationManager.java
+++ b/core/java/android/view/translation/TranslationManager.java
@@ -100,10 +100,6 @@ public final class TranslationManager {
 
     private final ITranslationManager mService;
 
-    @Nullable
-    @GuardedBy("mLock")
-    private ITranslationDirectManager mDirectServiceBinder;
-
     @NonNull
     @GuardedBy("mLock")
     private final SparseArray<Translator> mTranslators = new SparseArray<>();
@@ -131,11 +127,72 @@ public final class TranslationManager {
     /**
      * Creates an on-device Translator for natural language translation.
      *
+     * @param translationContext {@link TranslationContext} containing the specs for creating the
+     *                                                     Translator.
+     * @param executor Executor to run callback operations
+     * @param callback {@link Consumer} to receive the translator. A {@code null} value is returned
+     *                                 if the service could not create the translator.
+     */
+    public void createOnDeviceTranslator(@NonNull TranslationContext translationContext,
+            @NonNull @CallbackExecutor Executor executor, @NonNull Consumer<Translator> callback) {
+        Objects.requireNonNull(translationContext, "translationContext cannot be null");
+        Objects.requireNonNull(executor, "executor cannot be null");
+        Objects.requireNonNull(callback, "callback cannot be null");
+
+        synchronized (mLock) {
+            // TODO(b/176464808): Disallow multiple Translator now, it will throw
+            //  IllegalStateException. Need to discuss if we can allow multiple Translators.
+            if (mTranslatorIds.containsKey(translationContext)) {
+                executor.execute(() -> callback.accept(
+                        mTranslators.get(mTranslatorIds.get(translationContext))));
+                return;
+            }
+
+            int translatorId;
+            do {
+                translatorId = Math.abs(ID_GENERATOR.nextInt());
+            } while (translatorId == 0 || mTranslators.indexOfKey(translatorId) >= 0);
+            final int tId = translatorId;
+
+            new Translator(mContext, translationContext, translatorId, this, mHandler, mService,
+                    new Consumer<Translator>() {
+                        @Override
+                        public void accept(Translator translator) {
+                            if (translator == null) {
+                                final long token = Binder.clearCallingIdentity();
+                                try {
+                                    executor.execute(() -> callback.accept(null));
+                                } finally {
+                                    Binder.restoreCallingIdentity(token);
+                                }
+                                return;
+                            }
+
+                            mTranslators.put(tId, translator);
+                            mTranslatorIds.put(translationContext, tId);
+                            final long token = Binder.clearCallingIdentity();
+                            try {
+                                executor.execute(() -> callback.accept(translator));
+                            } finally {
+                                Binder.restoreCallingIdentity(token);
+                            }
+                        }
+                    });
+        }
+    }
+
+    /**
+     * Creates an on-device Translator for natural language translation.
+     *
      * <p><strong>NOTE: </strong>Call on a worker thread.
      *
+     * @deprecated use {@link #createOnDeviceTranslator(TranslationContext, Executor, Consumer)}
+     * instead.
+     *
      * @param translationContext {@link TranslationContext} containing the specs for creating the
      *                                                     Translator.
      */
+    @Deprecated
     @Nullable
     @WorkerThread
     public Translator createOnDeviceTranslator(@NonNull TranslationContext translationContext) {
diff --git a/core/java/android/view/translation/Translator.java b/core/java/android/view/translation/Translator.java
index 6037302a9445..b0d95b3b0a0f 100644
--- a/core/java/android/view/translation/Translator.java
+++ b/core/java/android/view/translation/Translator.java
@@ -101,10 +101,18 @@ public class Translator {
     public static final String EXTRA_SESSION_ID = "sessionId";
 
     static class ServiceBinderReceiver extends IResultReceiver.Stub {
+        // TODO: refactor how translator is instantiated after removing deprecated createTranslator.
         private final WeakReference<Translator> mTranslator;
         private final CountDownLatch mLatch = new CountDownLatch(1);
         private int mSessionId;
 
+        private Consumer<Translator> mCallback;
+
+        ServiceBinderReceiver(Translator translator, Consumer<Translator> callback) {
+            mTranslator = new WeakReference<>(translator);
+            mCallback = callback;
+        }
+
         ServiceBinderReceiver(Translator translator) {
             mTranslator = new WeakReference<>(translator);
         }
@@ -126,6 +134,9 @@ public class Translator {
         public void send(int resultCode, Bundle resultData) {
             if (resultCode == STATUS_SYNC_CALL_FAIL) {
                 mLatch.countDown();
+                if (mCallback != null) {
+                    mCallback.accept(null);
+                }
                 return;
             }
             mSessionId = resultData.getInt(EXTRA_SESSION_ID);
@@ -146,6 +157,9 @@ public class Translator {
             }
             translator.setServiceBinder(binder);
             mLatch.countDown();
+            if (mCallback != null) {
+                mCallback.accept(translator);
+            }
         }
 
         // TODO(b/176464808): maybe make SyncResultReceiver.TimeoutException constructor public
@@ -165,6 +179,32 @@ public class Translator {
     public Translator(@NonNull Context context,
             @NonNull TranslationContext translationContext, int sessionId,
             @NonNull TranslationManager translationManager, @NonNull Handler handler,
+            @Nullable ITranslationManager systemServerBinder,
+            @NonNull Consumer<Translator> callback) {
+        mContext = context;
+        mTranslationContext = translationContext;
+        mId = sessionId;
+        mManager = translationManager;
+        mHandler = handler;
+        mSystemServerBinder = systemServerBinder;
+        mServiceBinderReceiver = new ServiceBinderReceiver(this, callback);
+
+        try {
+            mSystemServerBinder.onSessionCreated(mTranslationContext, mId,
+                    mServiceBinderReceiver, mContext.getUserId());
+        } catch (RemoteException e) {
+            Log.w(TAG, "RemoteException calling startSession(): " + e);
+        }
+    }
+
+    /**
+     * Create the Translator.
+     *
+     * @hide
+     */
+    public Translator(@NonNull Context context,
+            @NonNull TranslationContext translationContext, int sessionId,
+            @NonNull TranslationManager translationManager, @NonNull Handler handler,
             @Nullable ITranslationManager systemServerBinder) {
         mContext = context;
         mTranslationContext = translationContext;