AIDL effect: add draining state support

 - effects can transit to DRAINING state with STOP command and continue
   the fade out processing
 - add draining in the audio eraser effect as an example
 - update VTS to support draining

Flag: EXEMPT bugfix
Bug: 379776482
Test: --test-mapping hardware/interfaces/audio/aidl/vts:presubmit

Change-Id: I2ca25cd085d1b6ae6cf8b0d1b58cd713aef0f7e5
Merged-In: I2ca25cd085d1b6ae6cf8b0d1b58cd713aef0f7e5
diff --git a/audio/aidl/aidl_api/android.hardware.audio.effect/current/android/hardware/audio/effect/State.aidl b/audio/aidl/aidl_api/android.hardware.audio.effect/current/android/hardware/audio/effect/State.aidl
index 17f9814..873fb43 100644
--- a/audio/aidl/aidl_api/android.hardware.audio.effect/current/android/hardware/audio/effect/State.aidl
+++ b/audio/aidl/aidl_api/android.hardware.audio.effect/current/android/hardware/audio/effect/State.aidl
@@ -37,4 +37,5 @@
   INIT,
   IDLE,
   PROCESSING,
+  DRAINING,
 }
diff --git a/audio/aidl/android/hardware/audio/effect/CommandId.aidl b/audio/aidl/android/hardware/audio/effect/CommandId.aidl
index d940b42..de573bf 100644
--- a/audio/aidl/android/hardware/audio/effect/CommandId.aidl
+++ b/audio/aidl/android/hardware/audio/effect/CommandId.aidl
@@ -33,25 +33,39 @@
     /**
      * Start effect engine processing.
      * An effect instance must start processing data and transfer to PROCESSING state if it is in
-     * IDLE state and have all necessary information. Otherwise it must:
-     * 1. Throw a EX_ILLEGAL_STATE exception if effect is not in IDLE state, or
-     * 2. Throw a EX_TRANSACTION_FAILED for all other errors.
+     * IDLE or DRAINING state and has all necessary information. Otherwise, it must:
+     * 1. Throw an EX_ILLEGAL_STATE exception if the effect is not in IDLE or DRAINING state, or
+     * 2. Throw an EX_TRANSACTION_FAILED for all other errors.
      *
-     * Depending on parameters set to the effect instance, effect may do process or reverse
-     * process after START command.
+     * If an effect instance in DRAINING state receives a START command, it must transit back to
+     * PROCESSING state.
      */
     START = 0,
     /**
-     * Stop effect engine processing with all resource kept.
-     * The currently processed audio data will be discarded if the effect engine is in PROCESSING
-     * state.
-     * Effect instance must do nothing and return ok when it receive STOP command in IDLE state.
+     * Stop effect engine processing with all resources kept.
+     * If the effect is in **PROCESSING** state:
+     *   - It must transition to **IDLE** state if no intermediate operations are required.
+     *   - It must transition to **DRAINING** state if draining (e.g., fading) is required.
+     *     - The instance must automatically transition to **IDLE** after draining.
+     *     - It must ignore any new `STOP` commands during **DRAINING**.
+     *     - `START` commands during **DRAINING** must transition the instance back to
+     *       **PROCESSING**.
+     * If the effect instance is already in **IDLE** state, it must do nothing and return success.
+     *
+     * If the effect instance transitions to DRAINING state:
+     * 1. It must automatically transition to IDLE after completing draining tasks.
+     * 2. It must ignore any new STOP commands received during the DRAINING state.
+     * 3. START commands during DRAINING must immediately transfer the instance back to PROCESSING.
+     *
      */
     STOP = 1,
     /**
      * Keep all parameter settings but reset the buffer content, stop engine processing, and transit
-     * instance state to IDLE if its in PROCESSING state.
+     * the instance state to IDLE if it is in PROCESSING state.
      * Effect instance must be able to handle RESET command at IDLE and PROCESSING states.
+     *
+     * If the implementation includes intermediate operations such as draining, the RESET command
+     * must bypass DRAINING and immediately transition the state to IDLE.
      */
     RESET = 2,
 
diff --git a/audio/aidl/android/hardware/audio/effect/State.aidl b/audio/aidl/android/hardware/audio/effect/State.aidl
index 85a4afc..1b698d7 100644
--- a/audio/aidl/android/hardware/audio/effect/State.aidl
+++ b/audio/aidl/android/hardware/audio/effect/State.aidl
@@ -24,18 +24,18 @@
  * it should transfer to IDLE state after handle the command successfully. Effect instance should
  * consume minimal resource and transfer to INIT state after it was close().
  *
- * Refer to State.gv for detailed state diagram.
+ * Refer to the state machine diagram `state.gv` for a detailed state diagram.
  */
 @VintfStability
 @Backing(type="byte")
 enum State {
-
     /**
      * An effect instance is in INIT state by default after it was created with
      * IFactory.createEffect(). When an effect instance is in INIT state, it should have instance
      * context initialized, and ready to handle IEffect.setParameter(), IEffect.open() as well as
      * all getter interfaces.
      *
+     * **Requirements in INIT state:**
      * In INIT state, effect instance must:
      * 1. Not handle any IEffect.command() and return EX_ILLEGAL_STATE with any Command.Id.
      * 2. Be able to handle all parameter setting with IEffect.setParameter().
@@ -43,28 +43,32 @@
      * IEffect.getState().
      * 4. Be able to handle IEffect.open() successfully after configuration.
      *
-     * Client is expected to do necessary configuration with IEffect.setParameter(), get all
-     * resource ready with IEffect.open(), and make sure effect instance transfer to IDLE state
-     * before sending commands with IEffect.command() interface. Effect instance must transfer
-     * from INIT to IDLE state after handle IEffect.open() call successfully.
+     * **State Transitions:**
+     * - Transitions to **IDLE** after successful `IEffect.open()`.
+     * - Remains in **INIT** on `IEffect.getState()` and `IEffect.getDescriptor()`.
+     * - Transitions to the final state on `IFactory.destroyEffect()`.
      */
     INIT,
+
     /**
      * An effect instance transfer to IDLE state after it was open successfully with IEffect.open()
      * in INIT state, or after it was stop/reset with Command.Id.STOP/RESET in PROCESSING state.
      *
-     * In IDLE state, effect instance must:
+     * **Requirements in IDLE state:**
      * 1. Be able to start effect processing engine with IEffect.command(Command.Id.START) call.
      * 2. Be able to handle all parameter setting with IEffect.setParameter().
      * 3. Be able to handle all getter interface calls like IEffect.getParameter() and
      * IEffect.getState().
      *
-     * The following state transfer can happen in IDLE state:
-     * 1. Transfer to PROCESSING if instance receive an START command and start processing data
-     * successfully.
-     * 2. Transfer to INIT if instance receive a close() call.
+     * **State Transitions:**
+     * - Transitions to **PROCESSING** on `IEffect.command(CommandId.START)` after starting
+     *   processing data successfully.
+     * - Transitions to **INIT** on `IEffect.close()`.
+     * - Remains in **IDLE** on `IEffect.getParameter()`, `IEffect.setParameter()`,
+     *   `IEffect.getDescriptor()`, `IEffect.command(CommandId.RESET)`, and `IEffect.reopen()`.
      */
     IDLE,
+
     /**
      * An effect instance is in PROCESSING state after it receive an START command and start
      * processing data successfully. Effect instance will transfer from PROCESSING to IDLE state if
@@ -75,12 +79,50 @@
      * the case of a close() call received when instance in PROCESSING state, it should try to stop
      * processing and transfer to IDLE first before close().
      *
-     * In PROCESSING state, effect instance must:
+     * **Requirements in PROCESSING state:**
      * 1. Return EX_ILLEGAL_STATE if it's not able to handle any parameter settings at runtime.
      * 2. Be able to handle STOP and RESET for IEffect.command() interface, and return
      * EX_ILLEGAL_STATE for all other commands.
      * 3. Must be able to handle all get* interface calls like IEffect.getParameter() and
      * IEffect.getState().
+     *
+     * **State Transitions:**
+     * - Transitions to **IDLE** on `IEffect.command(CommandId.STOP)` ( if no draining is required
+     *   or implemented) or `IEffect.command(CommandId.RESET)`.
+     * - Transitions to **DRAINING** on `IEffect.command(CommandId.STOP)` if draining is required.
+     * - Remains in **PROCESSING** on `IEffect.getParameter()`, `IEffect.setParameter()`,
+     *   `IEffect.getDescriptor()`, and `IEffect.reopen()`.
+     *
+     * **Notes:**
+     * - Clients should avoid calling `IEffect.close()` directly in this state; instead, they should
+     *   stop processing with `CommandId.STOP` before closing.
+     * - If `IEffect.close()` is called in this state, the effect instance should stop processing,
+     *   transition to **IDLE**, and then close.
      */
     PROCESSING,
+
+    /**
+     * DRAINING is an optional transitional state where the effect instance completes processing
+     * remaining input buffers or finalizes operations (e.g., fading) before stopping completely.
+     * This state is typically entered after a `CommandId.STOP` command in the PROCESSING state when
+     * draining is required.
+     *
+     * **Requirements in DRAINING state:**
+     * 1. Must handle `CommandId.START` and transition back to **PROCESSING**.
+     * 2. Must handle getter interface calls like `IEffect.getParameter()` and `IEffect.getState()`.
+     * 3. Must automatically transition to **IDLE** after draining is complete.
+     *
+     * **State Transitions:**
+     * - Transitions to **PROCESSING** on `IEffect.command(CommandId.START)`.
+     * - Transitions to **IDLE** on `IEffect.command(CommandId.RESET)`.
+     * - Transitions to **IDLE** automatically after draining is complete.
+     * - Remains in **DRAINING** on `IEffect.getParameter()`, `IEffect.setParameter()`,
+     *   `IEffect.getDescriptor()`, and `IEffect.reopen()`.
+     *
+     * **Notes:**
+     * - If not implemented, the effect instance may transition directly from **PROCESSING** to
+     *   **IDLE** without this intermediate state.
+     * - Any `CommandId.STOP` commands received during **DRAINING** should be ignored.
+     */
+    DRAINING,
 }
diff --git a/audio/aidl/android/hardware/audio/effect/state.gv b/audio/aidl/android/hardware/audio/effect/state.gv
index 22c70c8..2a8194e 100644
--- a/audio/aidl/android/hardware/audio/effect/state.gv
+++ b/audio/aidl/android/hardware/audio/effect/state.gv
@@ -13,26 +13,56 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-
 // To render: "dot -Tpng state.gv -o state.png"
+
 digraph effect_state_machine {
-    node[shape = point style = filled fillcolor = black width = 0.5] I;
-    node[shape = doublecircle] F;
-    node[shape = oval width = 1];
-    node[fillcolor = lightgreen] INIT;
-    node[fillcolor = lightblue] IDLE;
-    node[fillcolor = lightyellow] PROCESSING;
 
-    I -> INIT[label = "IFactory.createEffect" labelfontcolor = "navy"];
-    INIT -> F[label = "IFactory.destroyEffect"];
-    INIT -> IDLE[label = "IEffect.open()" labelfontcolor = "lime"];
-    IDLE -> PROCESSING[label = "IEffect.command(START"];
-    PROCESSING -> IDLE[label = "IEffect.command(STOP)\nIEffect.command(RESET)"];
-    IDLE -> INIT[label = "IEffect.close()"];
+    rankdir=LR; // Left to Right layout
 
-    INIT -> INIT[label = "IEffect.getState\nIEffect.getDescriptor"];
-    IDLE -> IDLE[label = "IEffect.getParameter\nIEffect.setParameter\nIEffect.getDescriptor\nIEffect.command(RESET)\nIEffect.reopen"];
-    PROCESSING
-            -> PROCESSING
-                    [label = "IEffect.getParameter\nIEffect.setParameter\nIEffect.getDescriptor\nIEffect.reopen"];
+    label="Effect State Machine";
+    fontsize=20;
+    labelloc=top;
+
+    node [fontname="Helvetica", fontsize=12, style=filled];
+
+    // Initial state node
+    I [shape=point, fillcolor=black, width=0.2];
+
+    // Final state node
+    F [shape=doublecircle, fillcolor=white, width=0.2];
+
+    // Define other nodes with colors
+    INIT [shape=ellipse, fillcolor=lightgreen];
+    IDLE [shape=ellipse, fillcolor=lightblue];
+    PROCESSING [shape=ellipse, fillcolor=lightyellow];
+    DRAINING [shape=ellipse, fillcolor=lightgrey];
+
+    // Transitions
+    I -> INIT [label="IFactory.createEffect", fontcolor="navy"];
+
+    INIT -> F [label="IFactory.destroyEffect"];
+
+    INIT -> IDLE [label="IEffect.open()", fontcolor="lime"];
+
+    IDLE -> PROCESSING [label="IEffect.command(START)"];
+
+    PROCESSING -> IDLE [label="IEffect.command(STOP)\nIEffect.command(RESET)"];
+
+    PROCESSING -> DRAINING [label="IEffect.command(STOP)", fontcolor="orange"];
+
+    DRAINING -> IDLE [label="Draining complete\n(IEffect.command(RESET)\nautomatic)"];
+
+    DRAINING -> PROCESSING [label="IEffect.command(START)\n(Interrupt draining)"];
+
+    IDLE -> INIT [label="IEffect.close()"];
+
+    // Self-loops
+    INIT -> INIT [label="IEffect.getState\nIEffect.getDescriptor"];
+
+    IDLE -> IDLE [label="IEffect.getParameter\nIEffect.setParameter\nIEffect.getDescriptor\nIEffect.command(RESET)\nIEffect.reopen"];
+
+    PROCESSING -> PROCESSING [label="IEffect.getParameter\nIEffect.setParameter\nIEffect.getDescriptor\nIEffect.reopen"];
+
+    DRAINING -> DRAINING [label="IEffect.getParameter\nIEffect.setParameter\nIEffect.getDescriptor\nIEffect.reopen\nFading"];
+
 }
diff --git a/audio/aidl/default/EffectContext.cpp b/audio/aidl/default/EffectContext.cpp
index 26c88b2..b354dd1 100644
--- a/audio/aidl/default/EffectContext.cpp
+++ b/audio/aidl/default/EffectContext.cpp
@@ -258,4 +258,18 @@
     return RetCode::SUCCESS;
 }
 
+RetCode EffectContext::startDraining() {
+    mIsDraining = true;
+    return RetCode::SUCCESS;
+}
+
+RetCode EffectContext::finishDraining() {
+    mIsDraining = false;
+    return RetCode::SUCCESS;
+}
+
+bool EffectContext::isDraining() {
+    return mIsDraining;
+}
+
 }  // namespace aidl::android::hardware::audio::effect
diff --git a/audio/aidl/default/EffectImpl.cpp b/audio/aidl/default/EffectImpl.cpp
index 3e61335..7857f53 100644
--- a/audio/aidl/default/EffectImpl.cpp
+++ b/audio/aidl/default/EffectImpl.cpp
@@ -79,7 +79,6 @@
     std::lock_guard lg(mImplMutex);
     RETURN_IF(mState == State::INIT, EX_ILLEGAL_STATE, "alreadyClosed");
 
-    // TODO: b/302036943 add reopen implementation
     RETURN_IF(!mImplContext, EX_NULL_POINTER, "nullContext");
     mImplContext->dupeFmqWithReopen(ret);
     return ndk::ScopedAStatus::ok();
@@ -347,7 +346,7 @@
 
     {
         std::lock_guard lg(mImplMutex);
-        if (mState != State::PROCESSING) {
+        if (mState != State::PROCESSING && mState != State::DRAINING) {
             LOG(DEBUG) << getEffectNameWithVersion()
                        << " skip process in state: " << toString(mState);
             return;
diff --git a/audio/aidl/default/EffectThread.cpp b/audio/aidl/default/EffectThread.cpp
index b515385..1a52c13 100644
--- a/audio/aidl/default/EffectThread.cpp
+++ b/audio/aidl/default/EffectThread.cpp
@@ -68,7 +68,11 @@
 RetCode EffectThread::startThread() {
     {
         std::lock_guard lg(mThreadMutex);
-        mStop = false;
+        if (mDraining) {
+            mDraining = false;
+        } else {
+            mStop = false;
+        }
         mCv.notify_one();
     }
 
@@ -87,6 +91,25 @@
     return RetCode::SUCCESS;
 }
 
+RetCode EffectThread::startDraining() {
+    std::lock_guard lg(mThreadMutex);
+    mDraining = true;
+    mCv.notify_one();
+
+    LOG(VERBOSE) << mName << __func__;
+    return RetCode::SUCCESS;
+}
+
+RetCode EffectThread::finishDraining() {
+    std::lock_guard lg(mThreadMutex);
+    mDraining = false;
+    mStop = true;
+    mCv.notify_one();
+
+    LOG(VERBOSE) << mName << __func__;
+    return RetCode::SUCCESS;
+}
+
 void EffectThread::threadLoop() {
     pthread_setname_np(pthread_self(), mName.substr(0, kMaxTaskNameLen - 1).c_str());
     setpriority(PRIO_PROCESS, 0, mPriority);
diff --git a/audio/aidl/default/audio_effects_config.xml b/audio/aidl/default/audio_effects_config.xml
index 2cef9eb..2e860d8 100644
--- a/audio/aidl/default/audio_effects_config.xml
+++ b/audio/aidl/default/audio_effects_config.xml
@@ -76,6 +76,7 @@
         <effect name="bassboost" library="bundle" uuid="8631f300-72e2-11df-b57e-0002a5d5c51b"/>
         <effect name="downmix" library="downmix" uuid="93f04452-e4fe-41cc-91f9-e475b6d1d69f"/>
         <effect name="dynamics_processing" library="dynamics_processing" uuid="e0e6539b-1781-7261-676f-6d7573696340"/>
+        <effect name="eraser" library="erasersw" uuid="fa81ab46-588b-11ed-9b6a-0242ac120002"/>
         <effect name="haptic_generator" library="haptic_generator" uuid="97c4acd1-8b82-4f2f-832e-c2fe5d7a9931"/>
         <effect name="loudness_enhancer" library="loudness_enhancer" uuid="fa415329-2034-4bea-b5dc-5b381c8d1e2c"/>
         <effect name="reverb_env_aux" library="reverb" uuid="4a387fc0-8ab3-11df-8bad-0002a5d5c51b"/>
diff --git a/audio/aidl/default/eraser/Eraser.cpp b/audio/aidl/default/eraser/Eraser.cpp
index 157ec79..59cc9a2 100644
--- a/audio/aidl/default/eraser/Eraser.cpp
+++ b/audio/aidl/default/eraser/Eraser.cpp
@@ -133,10 +133,70 @@
     LOG(DEBUG) << __func__;
 }
 
+ndk::ScopedAStatus EraserSw::command(CommandId command) {
+    std::lock_guard lg(mImplMutex);
+    RETURN_IF(mState == State::INIT, EX_ILLEGAL_STATE, "instanceNotOpen");
+
+    switch (command) {
+        case CommandId::START:
+            RETURN_OK_IF(mState == State::PROCESSING);
+            mState = State::PROCESSING;
+            mContext->enable();
+            startThread();
+            RETURN_IF(notifyEventFlag(mDataMqNotEmptyEf) != RetCode::SUCCESS, EX_ILLEGAL_STATE,
+                      "notifyEventFlagNotEmptyFailed");
+            break;
+        case CommandId::STOP:
+            RETURN_OK_IF(mState == State::IDLE || mState == State::DRAINING);
+            if (mVersion < kDrainSupportedVersion) {
+                mState = State::IDLE;
+                stopThread();
+                mContext->disable();
+            } else {
+                mState = State::DRAINING;
+                startDraining();
+                mContext->startDraining();
+            }
+            RETURN_IF(notifyEventFlag(mDataMqNotEmptyEf) != RetCode::SUCCESS, EX_ILLEGAL_STATE,
+                      "notifyEventFlagNotEmptyFailed");
+            break;
+        case CommandId::RESET:
+            mState = State::IDLE;
+            RETURN_IF(notifyEventFlag(mDataMqNotEmptyEf) != RetCode::SUCCESS, EX_ILLEGAL_STATE,
+                      "notifyEventFlagNotEmptyFailed");
+            stopThread();
+            mImplContext->disable();
+            mImplContext->reset();
+            mImplContext->resetBuffer();
+            break;
+        default:
+            LOG(ERROR) << getEffectNameWithVersion() << __func__ << " instance still processing";
+            return ndk::ScopedAStatus::fromExceptionCodeWithMessage(EX_ILLEGAL_ARGUMENT,
+                                                                    "CommandIdNotSupported");
+    }
+    LOG(VERBOSE) << getEffectNameWithVersion() << __func__
+                 << " transfer to state: " << toString(mState);
+    return ndk::ScopedAStatus::ok();
+}
+
 // Processing method running in EffectWorker thread.
 IEffect::Status EraserSw::effectProcessImpl(float* in, float* out, int samples) {
     RETURN_VALUE_IF(!mContext, (IEffect::Status{EX_NULL_POINTER, 0, 0}), "nullContext");
-    return mContext->process(in, out, samples);
+    IEffect::Status procStatus{STATUS_NOT_ENOUGH_DATA, 0, 0};
+    procStatus = mContext->process(in, out, samples);
+    if (mState == State::DRAINING && procStatus.status == STATUS_NOT_ENOUGH_DATA) {
+        drainingComplete_l();
+    }
+
+    return procStatus;
+}
+
+void EraserSw::drainingComplete_l() {
+    if (mState != State::DRAINING) return;
+
+    LOG(DEBUG) << getEffectNameWithVersion() << __func__;
+    finishDraining();
+    mState = State::IDLE;
 }
 
 EraserSwContext::EraserSwContext(int statusDepth, const Parameter::Common& common)
@@ -164,24 +224,47 @@
 
 IEffect::Status EraserSwContext::process(float* in, float* out, int samples) {
     LOG(DEBUG) << __func__ << " in " << in << " out " << out << " samples " << samples;
-    IEffect::Status status = {EX_ILLEGAL_ARGUMENT, 0, 0};
-
+    IEffect::Status procStatus = {EX_ILLEGAL_ARGUMENT, 0, 0};
     const auto inputChannelCount = getChannelCount(mCommon.input.base.channelMask);
     const auto outputChannelCount = getChannelCount(mCommon.output.base.channelMask);
     if (inputChannelCount < outputChannelCount) {
         LOG(ERROR) << __func__ << " invalid channel count, in: " << inputChannelCount
                    << " out: " << outputChannelCount;
-        return status;
+        return procStatus;
     }
 
-    int iFrames = samples / inputChannelCount;
+    if (samples <= 0 || 0 != samples % inputChannelCount) {
+        LOG(ERROR) << __func__ << " invalid samples: " << samples;
+        return procStatus;
+    }
+
+    const int iFrames = samples / inputChannelCount;
+    const float gainPerSample = 1.f / iFrames;
     for (int i = 0; i < iFrames; i++) {
-        std::memcpy(out, in, outputChannelCount);
+        if (isDraining()) {
+            const float gain = (iFrames - i - 1) * gainPerSample;
+            for (size_t c = 0; c < outputChannelCount; c++) {
+                out[c] = in[c] * gain;
+            }
+        } else {
+            std::memcpy(out, in, outputChannelCount * sizeof(float));
+        }
+
         in += inputChannelCount;
         out += outputChannelCount;
     }
-    return {STATUS_OK, static_cast<int32_t>(iFrames * inputChannelCount),
-            static_cast<int32_t>(iFrames * outputChannelCount)};
+
+    // drain for one cycle
+    if (isDraining()) {
+        procStatus.status = STATUS_NOT_ENOUGH_DATA;
+        finishDraining();
+    } else {
+        procStatus.status = STATUS_OK;
+    }
+    procStatus.fmqConsumed = static_cast<int32_t>(iFrames * inputChannelCount);
+    procStatus.fmqProduced = static_cast<int32_t>(iFrames * outputChannelCount);
+
+    return procStatus;
 }
 
 }  // namespace aidl::android::hardware::audio::effect
diff --git a/audio/aidl/default/eraser/Eraser.h b/audio/aidl/default/eraser/Eraser.h
index 0d4eb8f..7bf2f57 100644
--- a/audio/aidl/default/eraser/Eraser.h
+++ b/audio/aidl/default/eraser/Eraser.h
@@ -63,6 +63,9 @@
     IEffect::Status effectProcessImpl(float* in, float* out, int samples)
             REQUIRES(mImplMutex) final;
 
+    ndk::ScopedAStatus command(CommandId command) final;
+    void drainingComplete_l() REQUIRES(mImplMutex);
+
   private:
     static const std::vector<Range::SpatializerRange> kRanges;
     std::shared_ptr<EraserSwContext> mContext GUARDED_BY(mImplMutex);
diff --git a/audio/aidl/default/include/effect-impl/EffectContext.h b/audio/aidl/default/include/effect-impl/EffectContext.h
index 02a4caa..9e44349 100644
--- a/audio/aidl/default/include/effect-impl/EffectContext.h
+++ b/audio/aidl/default/include/effect-impl/EffectContext.h
@@ -86,7 +86,12 @@
     virtual RetCode disable();
     virtual RetCode reset();
 
+    virtual RetCode startDraining();
+    virtual RetCode finishDraining();
+    virtual bool isDraining();
+
   protected:
+    bool mIsDraining = false;
     int mVersion = 0;
     size_t mInputFrameSize = 0;
     size_t mOutputFrameSize = 0;
diff --git a/audio/aidl/default/include/effect-impl/EffectThread.h b/audio/aidl/default/include/effect-impl/EffectThread.h
index ec2a658..9abcdb8 100644
--- a/audio/aidl/default/include/effect-impl/EffectThread.h
+++ b/audio/aidl/default/include/effect-impl/EffectThread.h
@@ -38,6 +38,8 @@
     RetCode destroyThread();
     RetCode startThread();
     RetCode stopThread();
+    RetCode startDraining();
+    RetCode finishDraining();
 
     // Will call process() in a loop if the thread is running.
     void threadLoop();
@@ -49,6 +51,9 @@
      */
     virtual void process() = 0;
 
+  protected:
+    bool mDraining GUARDED_BY(mThreadMutex) = false;
+
   private:
     static constexpr int kMaxTaskNameLen = 15;
 
diff --git a/audio/aidl/vts/EffectHelper.h b/audio/aidl/vts/EffectHelper.h
index 570ecef..787bd12 100644
--- a/audio/aidl/vts/EffectHelper.h
+++ b/audio/aidl/vts/EffectHelper.h
@@ -124,7 +124,7 @@
             return;
         }
 
-        ASSERT_NO_FATAL_FAILURE(expectState(effect, State::IDLE));
+        ASSERT_TRUE(expectState(effect, State::IDLE));
         updateFrameSize(common);
     }
 
@@ -155,7 +155,7 @@
         if (effect) {
             ASSERT_STATUS(status, effect->close());
             if (status == EX_NONE) {
-                ASSERT_NO_FATAL_FAILURE(expectState(effect, State::INIT));
+                ASSERT_TRUE(expectState(effect, State::INIT));
             }
         }
     }
@@ -166,12 +166,14 @@
         ASSERT_STATUS(status, effect->getDescriptor(&desc));
     }
 
-    static void expectState(std::shared_ptr<IEffect> effect, State expectState,
-                            binder_status_t status = EX_NONE) {
-        ASSERT_NE(effect, nullptr);
-        State state;
-        ASSERT_STATUS(status, effect->getState(&state));
-        ASSERT_EQ(expectState, state);
+    static bool expectState(std::shared_ptr<IEffect> effect, State expectState) {
+        if (effect == nullptr) return false;
+
+        if (State state; EX_NONE != effect->getState(&state).getStatus() || expectState != state) {
+            return false;
+        }
+
+        return true;
     }
 
     static void commandIgnoreRet(std::shared_ptr<IEffect> effect, CommandId command) {
@@ -190,12 +192,14 @@
 
         switch (command) {
             case CommandId::START:
-                ASSERT_NO_FATAL_FAILURE(expectState(effect, State::PROCESSING));
+                ASSERT_TRUE(expectState(effect, State::PROCESSING));
                 break;
             case CommandId::STOP:
-                FALLTHROUGH_INTENDED;
+                ASSERT_TRUE(expectState(effect, State::IDLE) ||
+                            expectState(effect, State::DRAINING));
+                break;
             case CommandId::RESET:
-                ASSERT_NO_FATAL_FAILURE(expectState(effect, State::IDLE));
+                ASSERT_TRUE(expectState(effect, State::IDLE));
                 break;
             default:
                 return;
@@ -371,6 +375,24 @@
         return functor(result);
     }
 
+    // keep writing data to the FMQ until effect transit from DRAINING to IDLE
+    static void waitForDrain(std::vector<float>& inputBuffer, std::vector<float>& outputBuffer,
+                             const std::shared_ptr<IEffect>& effect,
+                             std::unique_ptr<EffectHelper::StatusMQ>& statusMQ,
+                             std::unique_ptr<EffectHelper::DataMQ>& inputMQ,
+                             std::unique_ptr<EffectHelper::DataMQ>& outputMQ, int version) {
+        State state;
+        while (effect->getState(&state).getStatus() == EX_NONE && state == State::DRAINING) {
+            EXPECT_NO_FATAL_FAILURE(
+                    EffectHelper::writeToFmq(statusMQ, inputMQ, inputBuffer, version));
+            EXPECT_NO_FATAL_FAILURE(EffectHelper::readFromFmq(
+                    statusMQ, 1, outputMQ, outputBuffer.size(), outputBuffer, std::nullopt));
+        }
+        ASSERT_TRUE(State::IDLE == state);
+        EXPECT_NO_FATAL_FAILURE(EffectHelper::readFromFmq(statusMQ, 0, outputMQ, 0, outputBuffer));
+        return;
+    }
+
     static void processAndWriteToOutput(std::vector<float>& inputBuffer,
                                         std::vector<float>& outputBuffer,
                                         const std::shared_ptr<IEffect>& effect,
@@ -404,8 +426,9 @@
         // Disable the process
         if (callStopReset) {
             ASSERT_NO_FATAL_FAILURE(command(effect, CommandId::STOP));
+            EXPECT_NO_FATAL_FAILURE(waitForDrain(inputBuffer, outputBuffer, effect, statusMQ,
+                                                 inputMQ, outputMQ, version));
         }
-        EXPECT_NO_FATAL_FAILURE(EffectHelper::readFromFmq(statusMQ, 0, outputMQ, 0, outputBuffer));
 
         if (callStopReset) {
             ASSERT_NO_FATAL_FAILURE(command(effect, CommandId::RESET));