audio: Allow Stream SM to stay in DRAINING state for early notify

Since for the "early notify" drain the stream sends `onDrainReady`
sligtly before the draining has actually finished, it should be
allowed to stay in the `DRAINING` state while sending this
notification. The stream gets into the `IDLE` state once the drain
has completed. While draining is ongoing, it is allowed for
the client to pause the stream.

As this is a new behavior in audio.core V3, extend VTS to skip
tests that verify this behavior for previous audio HAL versions.

Also, add previously missed `DrainOutAsync` scenario to the tests
list.

Bug: 363958142
Test: atest VtsHalAudioCoreTargetTest
Change-Id: I6878b07c2fa1ce4f19b3c9c5d0e0f2e5288e55c4
diff --git a/audio/aidl/default/Module.cpp b/audio/aidl/default/Module.cpp
index 51b6085..e96cf81 100644
--- a/audio/aidl/default/Module.cpp
+++ b/audio/aidl/default/Module.cpp
@@ -207,9 +207,9 @@
         return ndk::ScopedAStatus::fromExceptionCode(EX_ILLEGAL_ARGUMENT);
     }
     const auto& flags = portConfigIt->flags.value();
-    StreamContext::DebugParameters params{mDebug.streamTransientStateDelayMs,
-                                          mVendorDebug.forceTransientBurst,
-                                          mVendorDebug.forceSynchronousDrain};
+    StreamContext::DebugParameters params{
+            mDebug.streamTransientStateDelayMs, mVendorDebug.forceTransientBurst,
+            mVendorDebug.forceSynchronousDrain, mVendorDebug.forceDrainToDraining};
     std::unique_ptr<StreamContext::DataMQ> dataMQ = nullptr;
     std::shared_ptr<IStreamCallback> streamAsyncCallback = nullptr;
     std::shared_ptr<ISoundDose> soundDose;
@@ -1524,6 +1524,7 @@
 
 const std::string Module::VendorDebug::kForceTransientBurstName = "aosp.forceTransientBurst";
 const std::string Module::VendorDebug::kForceSynchronousDrainName = "aosp.forceSynchronousDrain";
+const std::string Module::VendorDebug::kForceDrainToDrainingName = "aosp.forceDrainToDraining";
 
 ndk::ScopedAStatus Module::getVendorParameters(const std::vector<std::string>& in_ids,
                                                std::vector<VendorParameter>* _aidl_return) {
@@ -1538,6 +1539,10 @@
             VendorParameter forceSynchronousDrain{.id = id};
             forceSynchronousDrain.ext.setParcelable(Boolean{mVendorDebug.forceSynchronousDrain});
             _aidl_return->push_back(std::move(forceSynchronousDrain));
+        } else if (id == VendorDebug::kForceDrainToDrainingName) {
+            VendorParameter forceDrainToDraining{.id = id};
+            forceDrainToDraining.ext.setParcelable(Boolean{mVendorDebug.forceDrainToDraining});
+            _aidl_return->push_back(std::move(forceDrainToDraining));
         } else {
             allParametersKnown = false;
             LOG(VERBOSE) << __func__ << ": " << mType << ": unrecognized parameter \"" << id << "\"";
@@ -1578,6 +1583,10 @@
             if (!extractParameter<Boolean>(p, &mVendorDebug.forceSynchronousDrain)) {
                 return ndk::ScopedAStatus::fromExceptionCode(EX_ILLEGAL_ARGUMENT);
             }
+        } else if (p.id == VendorDebug::kForceDrainToDrainingName) {
+            if (!extractParameter<Boolean>(p, &mVendorDebug.forceDrainToDraining)) {
+                return ndk::ScopedAStatus::fromExceptionCode(EX_ILLEGAL_ARGUMENT);
+            }
         } else {
             allParametersKnown = false;
             LOG(VERBOSE) << __func__ << ": " << mType << ": unrecognized parameter \"" << p.id
diff --git a/audio/aidl/default/Stream.cpp b/audio/aidl/default/Stream.cpp
index 3d7f30c..4525f6a 100644
--- a/audio/aidl/default/Stream.cpp
+++ b/audio/aidl/default/Stream.cpp
@@ -382,8 +382,20 @@
 const std::string StreamOutWorkerLogic::kThreadName = "writer";
 
 StreamOutWorkerLogic::Status StreamOutWorkerLogic::cycle() {
-    if (mState == StreamDescriptor::State::DRAINING ||
-        mState == StreamDescriptor::State::TRANSFERRING) {
+    if (mState == StreamDescriptor::State::DRAINING && mContext->getForceDrainToDraining() &&
+        mOnDrainReadyStatus == OnDrainReadyStatus::UNSENT) {
+        std::shared_ptr<IStreamCallback> asyncCallback = mContext->getAsyncCallback();
+        if (asyncCallback != nullptr) {
+            ndk::ScopedAStatus status = asyncCallback->onDrainReady();
+            if (!status.isOk()) {
+                LOG(ERROR) << __func__ << ": error from onDrainReady: " << status;
+            }
+            // This sets the timeout for moving into IDLE on next iterations.
+            switchToTransientState(StreamDescriptor::State::DRAINING);
+            mOnDrainReadyStatus = OnDrainReadyStatus::SENT;
+        }
+    } else if (mState == StreamDescriptor::State::DRAINING ||
+               mState == StreamDescriptor::State::TRANSFERRING) {
         if (auto stateDurationMs = std::chrono::duration_cast<std::chrono::milliseconds>(
                     std::chrono::steady_clock::now() - mTransientStateStart);
             stateDurationMs >= mTransientStateDelayMs) {
@@ -396,9 +408,12 @@
                 // drain or transfer completion. In the stub, we switch unconditionally.
                 if (mState == StreamDescriptor::State::DRAINING) {
                     mState = StreamDescriptor::State::IDLE;
-                    ndk::ScopedAStatus status = asyncCallback->onDrainReady();
-                    if (!status.isOk()) {
-                        LOG(ERROR) << __func__ << ": error from onDrainReady: " << status;
+                    if (mOnDrainReadyStatus != OnDrainReadyStatus::SENT) {
+                        ndk::ScopedAStatus status = asyncCallback->onDrainReady();
+                        if (!status.isOk()) {
+                            LOG(ERROR) << __func__ << ": error from onDrainReady: " << status;
+                        }
+                        mOnDrainReadyStatus = OnDrainReadyStatus::SENT;
                     }
                 } else {
                     mState = StreamDescriptor::State::ACTIVE;
@@ -537,6 +552,10 @@
                             mState = StreamDescriptor::State::IDLE;
                         } else {
                             switchToTransientState(StreamDescriptor::State::DRAINING);
+                            mOnDrainReadyStatus =
+                                    mode == StreamDescriptor::DrainMode::DRAIN_EARLY_NOTIFY
+                                            ? OnDrainReadyStatus::UNSENT
+                                            : OnDrainReadyStatus::IGNORE;
                         }
                     } else {
                         LOG(ERROR) << __func__ << ": drain failed: " << status;
diff --git a/audio/aidl/default/include/core-impl/Module.h b/audio/aidl/default/include/core-impl/Module.h
index 7e32cf2..d03598a 100644
--- a/audio/aidl/default/include/core-impl/Module.h
+++ b/audio/aidl/default/include/core-impl/Module.h
@@ -148,8 +148,10 @@
     struct VendorDebug {
         static const std::string kForceTransientBurstName;
         static const std::string kForceSynchronousDrainName;
+        static const std::string kForceDrainToDrainingName;
         bool forceTransientBurst = false;
         bool forceSynchronousDrain = false;
+        bool forceDrainToDraining = false;
     };
     // ids of device ports created at runtime via 'connectExternalDevice'.
     // Also stores a list of ids of mix ports with dynamic profiles that were populated from
diff --git a/audio/aidl/default/include/core-impl/Stream.h b/audio/aidl/default/include/core-impl/Stream.h
index f7b9269..8297fc5 100644
--- a/audio/aidl/default/include/core-impl/Stream.h
+++ b/audio/aidl/default/include/core-impl/Stream.h
@@ -78,6 +78,10 @@
         bool forceTransientBurst = false;
         // Force the "drain" command to be synchronous, going directly to the IDLE state.
         bool forceSynchronousDrain = false;
+        // Force the "drain early notify" command to keep the SM in the DRAINING state
+        // after sending 'onDrainReady' callback. The SM moves to IDLE after
+        // 'transientStateDelayMs'.
+        bool forceDrainToDraining = false;
     };
 
     StreamContext() = default;
@@ -119,6 +123,7 @@
     ::aidl::android::media::audio::common::AudioIoFlags getFlags() const { return mFlags; }
     bool getForceTransientBurst() const { return mDebugParameters.forceTransientBurst; }
     bool getForceSynchronousDrain() const { return mDebugParameters.forceSynchronousDrain; }
+    bool getForceDrainToDraining() const { return mDebugParameters.forceDrainToDraining; }
     size_t getFrameSize() const;
     int getInternalCommandCookie() const { return mInternalCommandCookie; }
     int32_t getMixPortHandle() const { return mMixPortHandle; }
@@ -301,6 +306,9 @@
     bool write(size_t clientSize, StreamDescriptor::Reply* reply);
 
     std::shared_ptr<IStreamOutEventCallback> mEventCallback;
+
+    enum OnDrainReadyStatus : int32_t { IGNORE /*used for DRAIN_ALL*/, UNSENT, SENT };
+    OnDrainReadyStatus mOnDrainReadyStatus = OnDrainReadyStatus::IGNORE;
 };
 using StreamOutWorker = StreamWorkerImpl<StreamOutWorkerLogic>;