Add more methods for capture and injection

Add #close() but keep it and the constructor as package-private because we want caller to call VirtualAudioDevice#close().

Add #getFormat() and the audio format will be owned by
AudioCapture/AudioInjection, not VirtualAudioSession.

Add all #read() and #write() overloads method, so that the caller can
replace the call to AudioRecord/AudioTrack easily.

Bug: 218542209
CTS-Coverage-Bug: 218528439
Test: atest FrameworksCoreTests:android.companion.virtual
Change-Id: Iccfd497e6aaeacc60a57d28ae730fdc0cce46a25

3 files changed, 204 insertions, 30 deletions

diff --git a/core/java/android/companion/virtual/audio/AudioCapture.java b/core/java/android/companion/virtual/audio/AudioCapture.java
index ebe17dba5775..d6d0d2b79c83 100644
--- a/core/java/android/companion/virtual/audio/AudioCapture.java
+++ b/core/java/android/companion/virtual/audio/AudioCapture.java
@@ -16,13 +16,16 @@
 
 package android.companion.virtual.audio;
 
+import static android.media.AudioRecord.READ_BLOCKING;
 import static android.media.AudioRecord.RECORDSTATE_RECORDING;
 import static android.media.AudioRecord.RECORDSTATE_STOPPED;
+import static android.media.AudioRecord.STATE_INITIALIZED;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.SuppressLint;
 import android.annotation.SystemApi;
+import android.media.AudioFormat;
 import android.media.AudioRecord;
 import android.util.Log;
 
@@ -42,12 +45,12 @@ import java.nio.ByteBuffer;
 public final class AudioCapture {
     private static final String TAG = "AudioCapture";
 
+    private final AudioFormat mAudioFormat;
     private final Object mLock = new Object();
 
     @GuardedBy("mLock")
     @Nullable
     private AudioRecord mAudioRecord;
-
     @GuardedBy("mLock")
     private int mRecordingState = RECORDSTATE_STOPPED;
 
@@ -63,12 +66,12 @@ public final class AudioCapture {
     void setAudioRecord(@Nullable AudioRecord audioRecord) {
         Log.d(TAG, "set AudioRecord with " + audioRecord);
         synchronized (mLock) {
-            // Release old reference.
-            if (mAudioRecord != null) {
-                mAudioRecord.release();
-            }
             // Sync recording state for new reference.
             if (audioRecord != null) {
+                if (audioRecord.getState() != STATE_INITIALIZED) {
+                    throw new IllegalStateException("set an uninitialized AudioRecord.");
+                }
+
                 if (mRecordingState == RECORDSTATE_RECORDING
                         && audioRecord.getRecordingState() != RECORDSTATE_RECORDING) {
                     audioRecord.startRecording();
@@ -78,16 +81,97 @@ public final class AudioCapture {
                     audioRecord.stop();
                 }
             }
+
+            // Release old reference before assigning the new reference.
+            if (mAudioRecord != null) {
+                mAudioRecord.release();
+            }
             mAudioRecord = audioRecord;
         }
     }
 
+    AudioCapture(@NonNull AudioFormat audioFormat) {
+        mAudioFormat = audioFormat;
+    }
+
+    void close() {
+        synchronized (mLock) {
+            if (mAudioRecord != null) {
+                mAudioRecord.release();
+                mAudioRecord = null;
+            }
+        }
+    }
+
+    /** See {@link AudioRecord#getFormat()} */
+    public @NonNull AudioFormat getFormat() {
+        return mAudioFormat;
+    }
+
+    /** See {@link AudioRecord#read(byte[], int, int)} */
+    public int read(@NonNull byte[] audioData, int offsetInBytes, int sizeInBytes) {
+        return read(audioData, offsetInBytes, sizeInBytes, READ_BLOCKING);
+    }
+
+    /** See {@link AudioRecord#read(byte[], int, int, int)} */
+    public int read(@NonNull byte[] audioData, int offsetInBytes, int sizeInBytes,
+            @AudioRecord.ReadMode int readMode) {
+        final int sizeRead;
+        synchronized (mLock) {
+            if (mAudioRecord != null) {
+                sizeRead = mAudioRecord.read(audioData, offsetInBytes, sizeInBytes, readMode);
+            } else {
+                sizeRead = 0;
+            }
+        }
+        return sizeRead;
+    }
+
     /** See {@link AudioRecord#read(ByteBuffer, int)}. */
     public int read(@NonNull ByteBuffer audioBuffer, int sizeInBytes) {
+        return read(audioBuffer, sizeInBytes, READ_BLOCKING);
+    }
+
+    /** See {@link AudioRecord#read(ByteBuffer, int, int)}. */
+    public int read(@NonNull ByteBuffer audioBuffer, int sizeInBytes,
+            @AudioRecord.ReadMode int readMode) {
+        final int sizeRead;
+        synchronized (mLock) {
+            if (mAudioRecord != null) {
+                sizeRead = mAudioRecord.read(audioBuffer, sizeInBytes, readMode);
+            } else {
+                sizeRead = 0;
+            }
+        }
+        return sizeRead;
+    }
+
+    /** See {@link AudioRecord#read(float[], int, int, int)}. */
+    public int read(@NonNull float[] audioData, int offsetInFloats, int sizeInFloats,
+            @AudioRecord.ReadMode int readMode) {
+        final int sizeRead;
+        synchronized (mLock) {
+            if (mAudioRecord != null) {
+                sizeRead = mAudioRecord.read(audioData, offsetInFloats, sizeInFloats, readMode);
+            } else {
+                sizeRead = 0;
+            }
+        }
+        return sizeRead;
+    }
+
+    /** See {@link AudioRecord#read(short[], int, int)}. */
+    public int read(@NonNull short[] audioData, int offsetInShorts, int sizeInShorts) {
+        return read(audioData, offsetInShorts, sizeInShorts, READ_BLOCKING);
+    }
+
+    /** See {@link AudioRecord#read(short[], int, int, int)}. */
+    public int read(@NonNull short[] audioData, int offsetInShorts, int sizeInShorts,
+            @AudioRecord.ReadMode int readMode) {
         final int sizeRead;
         synchronized (mLock) {
             if (mAudioRecord != null) {
-                sizeRead = mAudioRecord.read(audioBuffer, sizeInBytes);
+                sizeRead = mAudioRecord.read(audioData, offsetInShorts, sizeInShorts, readMode);
             } else {
                 sizeRead = 0;
             }
diff --git a/core/java/android/companion/virtual/audio/AudioInjection.java b/core/java/android/companion/virtual/audio/AudioInjection.java
index 5e8e0a487a2e..9d6a3eb84351 100644
--- a/core/java/android/companion/virtual/audio/AudioInjection.java
+++ b/core/java/android/companion/virtual/audio/AudioInjection.java
@@ -18,11 +18,14 @@ package android.companion.virtual.audio;
 
 import static android.media.AudioTrack.PLAYSTATE_PLAYING;
 import static android.media.AudioTrack.PLAYSTATE_STOPPED;
+import static android.media.AudioTrack.STATE_INITIALIZED;
+import static android.media.AudioTrack.WRITE_BLOCKING;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.SuppressLint;
 import android.annotation.SystemApi;
+import android.media.AudioFormat;
 import android.media.AudioTrack;
 import android.util.Log;
 
@@ -42,7 +45,9 @@ import java.nio.ByteBuffer;
 public final class AudioInjection {
     private static final String TAG = "AudioInjection";
 
+    private final AudioFormat mAudioFormat;
     private final Object mLock = new Object();
+
     @GuardedBy("mLock")
     @Nullable
     private AudioTrack mAudioTrack;
@@ -70,12 +75,12 @@ public final class AudioInjection {
     void setAudioTrack(@Nullable AudioTrack audioTrack) {
         Log.d(TAG, "set AudioTrack with " + audioTrack);
         synchronized (mLock) {
-            // Release old reference.
-            if (mAudioTrack != null) {
-                mAudioTrack.release();
-            }
             // Sync play state for new reference.
             if (audioTrack != null) {
+                if (audioTrack.getState() != STATE_INITIALIZED) {
+                    throw new IllegalStateException("set an uninitialized AudioTrack.");
+                }
+
                 if (mPlayState == PLAYSTATE_PLAYING
                         && audioTrack.getPlayState() != PLAYSTATE_PLAYING) {
                     audioTrack.play();
@@ -85,10 +90,52 @@ public final class AudioInjection {
                     audioTrack.stop();
                 }
             }
+
+            // Release old reference before assigning the new reference.
+            if (mAudioTrack != null) {
+                mAudioTrack.release();
+            }
             mAudioTrack = audioTrack;
         }
     }
 
+    AudioInjection(@NonNull AudioFormat audioFormat) {
+        mAudioFormat = audioFormat;
+    }
+
+    void close() {
+        synchronized (mLock) {
+            if (mAudioTrack != null) {
+                mAudioTrack.release();
+                mAudioTrack = null;
+            }
+        }
+    }
+
+    /** See {@link AudioTrack#getFormat()}. */
+    public @NonNull AudioFormat getFormat() {
+        return mAudioFormat;
+    }
+
+    /** See {@link AudioTrack#write(byte[], int, int)}. */
+    public int write(@NonNull byte[] audioData, int offsetInBytes, int sizeInBytes) {
+        return write(audioData, offsetInBytes, sizeInBytes, WRITE_BLOCKING);
+    }
+
+    /** See {@link AudioTrack#write(byte[], int, int, int)}. */
+    public int write(@NonNull byte[] audioData, int offsetInBytes, int sizeInBytes,
+            @AudioTrack.WriteMode int writeMode) {
+        final int sizeWrite;
+        synchronized (mLock) {
+            if (mAudioTrack != null && !mIsSilent) {
+                sizeWrite = mAudioTrack.write(audioData, offsetInBytes, sizeInBytes, writeMode);
+            } else {
+                sizeWrite = 0;
+            }
+        }
+        return sizeWrite;
+    }
+
     /** See {@link AudioTrack#write(ByteBuffer, int, int)}. */
     public int write(@NonNull ByteBuffer audioBuffer, int sizeInBytes, int writeMode) {
         final int sizeWrite;
@@ -102,6 +149,53 @@ public final class AudioInjection {
         return sizeWrite;
     }
 
+    /** See {@link AudioTrack#write(ByteBuffer, int, int, long)}. */
+    public int write(@NonNull ByteBuffer audioBuffer, int sizeInBytes,
+            @AudioTrack.WriteMode int writeMode, long timestamp) {
+        final int sizeWrite;
+        synchronized (mLock) {
+            if (mAudioTrack != null && !mIsSilent) {
+                sizeWrite = mAudioTrack.write(audioBuffer, sizeInBytes, writeMode, timestamp);
+            } else {
+                sizeWrite = 0;
+            }
+        }
+        return sizeWrite;
+    }
+
+    /** See {@link AudioTrack#write(float[], int, int, int)}. */
+    public int write(@NonNull float[] audioData, int offsetInFloats, int sizeInFloats,
+            @AudioTrack.WriteMode int writeMode) {
+        final int sizeWrite;
+        synchronized (mLock) {
+            if (mAudioTrack != null && !mIsSilent) {
+                sizeWrite = mAudioTrack.write(audioData, offsetInFloats, sizeInFloats, writeMode);
+            } else {
+                sizeWrite = 0;
+            }
+        }
+        return sizeWrite;
+    }
+
+    /** See {@link AudioTrack#write(short[], int, int)}. */
+    public int write(@NonNull short[] audioData, int offsetInShorts, int sizeInShorts) {
+        return write(audioData, offsetInShorts, sizeInShorts, WRITE_BLOCKING);
+    }
+
+    /** See {@link AudioTrack#write(short[], int, int, int)}. */
+    public int write(@NonNull short[] audioData, int offsetInShorts, int sizeInShorts,
+            @AudioTrack.WriteMode int writeMode) {
+        final int sizeWrite;
+        synchronized (mLock) {
+            if (mAudioTrack != null && !mIsSilent) {
+                sizeWrite = mAudioTrack.write(audioData, offsetInShorts, sizeInShorts, writeMode);
+            } else {
+                sizeWrite = 0;
+            }
+        }
+        return sizeWrite;
+    }
+
     /** See {@link AudioTrack#play()}. */
     public void play() {
         synchronized (mLock) {
diff --git a/core/java/android/companion/virtual/audio/VirtualAudioSession.java b/core/java/android/companion/virtual/audio/VirtualAudioSession.java
index c6a10456c331..524f6e2ab063 100644
--- a/core/java/android/companion/virtual/audio/VirtualAudioSession.java
+++ b/core/java/android/companion/virtual/audio/VirtualAudioSession.java
@@ -68,12 +68,6 @@ public final class VirtualAudioSession extends IAudioSessionCallback.Stub implem
     private AudioPolicy mAudioPolicy;
     @Nullable
     @GuardedBy("mLock")
-    private AudioFormat mCaptureFormat;
-    @Nullable
-    @GuardedBy("mLock")
-    private AudioFormat mInjectionFormat;
-    @Nullable
-    @GuardedBy("mLock")
     private AudioCapture mAudioCapture;
     @Nullable
     @GuardedBy("mLock")
@@ -104,8 +98,7 @@ public final class VirtualAudioSession extends IAudioSessionCallback.Stub implem
                         "Cannot start capture while another capture is ongoing.");
             }
 
-            mCaptureFormat = captureFormat;
-            mAudioCapture = new AudioCapture();
+            mAudioCapture = new AudioCapture(captureFormat);
             mAudioCapture.startRecording();
             return mAudioCapture;
         }
@@ -127,8 +120,7 @@ public final class VirtualAudioSession extends IAudioSessionCallback.Stub implem
                         "Cannot start injection while injection is already ongoing.");
             }
 
-            mInjectionFormat = injectionFormat;
-            mAudioInjection = new AudioInjection();
+            mAudioInjection = new AudioInjection(injectionFormat);
             mAudioInjection.play();
 
             mUserRestrictionsDetector.register(/* callback= */ this);
@@ -179,10 +171,14 @@ public final class VirtualAudioSession extends IAudioSessionCallback.Stub implem
         mUserRestrictionsDetector.unregister();
         releaseAudioStreams();
         synchronized (mLock) {
-            mAudioCapture = null;
-            mAudioInjection = null;
-            mCaptureFormat = null;
-            mInjectionFormat = null;
+            if (mAudioCapture != null) {
+                mAudioCapture.close();
+                mAudioCapture = null;
+            }
+            if (mAudioInjection != null) {
+                mAudioInjection.close();
+                mAudioInjection = null;
+            }
         }
     }
 
@@ -198,9 +194,9 @@ public final class VirtualAudioSession extends IAudioSessionCallback.Stub implem
     @RequiresPermission(android.Manifest.permission.MODIFY_AUDIO_ROUTING)
     private void createAudioStreams(int[] appUids) {
         synchronized (mLock) {
-            if (mCaptureFormat == null && mInjectionFormat == null) {
+            if (mAudioCapture == null && mAudioInjection == null) {
                 throw new IllegalStateException(
-                        "At least one of captureFormat and injectionFormat must be specified.");
+                        "At least one of AudioCapture and AudioInjection must be started.");
             }
             if (mAudioPolicy != null) {
                 throw new IllegalStateException(
@@ -218,12 +214,12 @@ public final class VirtualAudioSession extends IAudioSessionCallback.Stub implem
             AudioMix audioRecordMix = null;
             AudioMix audioTrackMix = null;
             AudioPolicy.Builder builder = new AudioPolicy.Builder(mContext);
-            if (mCaptureFormat != null) {
-                audioRecordMix = createAudioRecordMix(mCaptureFormat, appUids);
+            if (mAudioCapture != null) {
+                audioRecordMix = createAudioRecordMix(mAudioCapture.getFormat(), appUids);
                 builder.addMix(audioRecordMix);
             }
-            if (mInjectionFormat != null) {
-                audioTrackMix = createAudioTrackMix(mInjectionFormat, appUids);
+            if (mAudioInjection != null) {
+                audioTrackMix = createAudioTrackMix(mAudioInjection.getFormat(), appUids);
                 builder.addMix(audioTrackMix);
             }
             mAudioPolicy = builder.build();