Merge "Fix VibrationThread handling old HAL callbacks" into main

20 files changed, 338 insertions, 261 deletions

diff --git a/services/core/java/com/android/server/vibrator/AbstractVibratorStep.java b/services/core/java/com/android/server/vibrator/AbstractVibratorStep.java
index 07d9ad16aca5..da6478bd3395 100644
--- a/services/core/java/com/android/server/vibrator/AbstractVibratorStep.java
+++ b/services/core/java/com/android/server/vibrator/AbstractVibratorStep.java
@@ -124,6 +124,8 @@ abstract class AbstractVibratorStep extends Step {
             Slog.d(VibrationThread.TAG,
                     "Turning off vibrator " + getVibratorId());
         }
+        // Make sure we ignore any pending callback from old vibration commands.
+        conductor.nextVibratorCallbackStepId(getVibratorId());
         controller.off();
         getVibration().stats.reportVibratorOff();
         mPendingVibratorOffDeadline = 0;
diff --git a/services/core/java/com/android/server/vibrator/ComposePrimitivesVibratorStep.java b/services/core/java/com/android/server/vibrator/ComposePrimitivesVibratorStep.java
index e495af59a2f9..b3eead109999 100644
--- a/services/core/java/com/android/server/vibrator/ComposePrimitivesVibratorStep.java
+++ b/services/core/java/com/android/server/vibrator/ComposePrimitivesVibratorStep.java
@@ -73,7 +73,8 @@ final class ComposePrimitivesVibratorStep extends AbstractComposedVibratorStep {
 
             PrimitiveSegment[] primitivesArray =
                     primitives.toArray(new PrimitiveSegment[primitives.size()]);
-            long vibratorOnResult = controller.on(primitivesArray, getVibration().id);
+            int stepId = conductor.nextVibratorCallbackStepId(getVibratorId());
+            long vibratorOnResult = controller.on(primitivesArray, getVibration().id, stepId);
             handleVibratorOnResult(vibratorOnResult);
             getVibration().stats.reportComposePrimitives(vibratorOnResult, primitivesArray);
 
diff --git a/services/core/java/com/android/server/vibrator/ComposePwleV2VibratorStep.java b/services/core/java/com/android/server/vibrator/ComposePwleV2VibratorStep.java
index bb8e6eed5707..7b41457b5016 100644
--- a/services/core/java/com/android/server/vibrator/ComposePwleV2VibratorStep.java
+++ b/services/core/java/com/android/server/vibrator/ComposePwleV2VibratorStep.java
@@ -72,7 +72,8 @@ final class ComposePwleV2VibratorStep extends AbstractComposedVibratorStep {
                         + controller.getVibratorInfo().getId());
             }
             PwlePoint[] pwlesArray = pwles.toArray(new PwlePoint[pwles.size()]);
-            long vibratorOnResult = controller.on(pwlesArray, getVibration().id);
+            int stepId = conductor.nextVibratorCallbackStepId(getVibratorId());
+            long vibratorOnResult = controller.on(pwlesArray, getVibration().id, stepId);
             handleVibratorOnResult(vibratorOnResult);
             getVibration().stats.reportComposePwle(vibratorOnResult, pwlesArray);
 
diff --git a/services/core/java/com/android/server/vibrator/ComposePwleVibratorStep.java b/services/core/java/com/android/server/vibrator/ComposePwleVibratorStep.java
index e8952fafaf77..d003251ef307 100644
--- a/services/core/java/com/android/server/vibrator/ComposePwleVibratorStep.java
+++ b/services/core/java/com/android/server/vibrator/ComposePwleVibratorStep.java
@@ -71,7 +71,8 @@ final class ComposePwleVibratorStep extends AbstractComposedVibratorStep {
                         + controller.getVibratorInfo().getId());
             }
             RampSegment[] pwlesArray = pwles.toArray(new RampSegment[pwles.size()]);
-            long vibratorOnResult = controller.on(pwlesArray, getVibration().id);
+            int stepId = conductor.nextVibratorCallbackStepId(getVibratorId());
+            long vibratorOnResult = controller.on(pwlesArray, getVibration().id, stepId);
             handleVibratorOnResult(vibratorOnResult);
             getVibration().stats.reportComposePwle(vibratorOnResult, pwlesArray);
 
diff --git a/services/core/java/com/android/server/vibrator/ExternalVibrationSession.java b/services/core/java/com/android/server/vibrator/ExternalVibrationSession.java
index a92ac679b0f4..b2cc1f60ca1d 100644
--- a/services/core/java/com/android/server/vibrator/ExternalVibrationSession.java
+++ b/services/core/java/com/android/server/vibrator/ExternalVibrationSession.java
@@ -148,7 +148,7 @@ final class ExternalVibrationSession extends Vibration
     }
 
     @Override
-    public void notifyVibratorCallback(int vibratorId, long vibrationId) {
+    public void notifyVibratorCallback(int vibratorId, long vibrationId, long stepId) {
         // ignored, external control does not expect callbacks from the vibrator
     }
 
diff --git a/services/core/java/com/android/server/vibrator/PerformPrebakedVibratorStep.java b/services/core/java/com/android/server/vibrator/PerformPrebakedVibratorStep.java
index 4b23216258af..88bb181781bf 100644
--- a/services/core/java/com/android/server/vibrator/PerformPrebakedVibratorStep.java
+++ b/services/core/java/com/android/server/vibrator/PerformPrebakedVibratorStep.java
@@ -64,7 +64,8 @@ final class PerformPrebakedVibratorStep extends AbstractComposedVibratorStep {
             }
 
             VibrationEffect fallback = getVibration().getFallback(prebaked.getEffectId());
-            long vibratorOnResult = controller.on(prebaked, getVibration().id);
+            int stepId = conductor.nextVibratorCallbackStepId(getVibratorId());
+            long vibratorOnResult = controller.on(prebaked, getVibration().id, stepId);
             handleVibratorOnResult(vibratorOnResult);
             getVibration().stats.reportPerformEffect(vibratorOnResult, prebaked);
 
diff --git a/services/core/java/com/android/server/vibrator/PerformVendorEffectVibratorStep.java b/services/core/java/com/android/server/vibrator/PerformVendorEffectVibratorStep.java
index 407f3d996798..d4bcc36aee18 100644
--- a/services/core/java/com/android/server/vibrator/PerformVendorEffectVibratorStep.java
+++ b/services/core/java/com/android/server/vibrator/PerformVendorEffectVibratorStep.java
@@ -49,7 +49,8 @@ final class PerformVendorEffectVibratorStep extends AbstractVibratorStep {
     public List<Step> play() {
         Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "PerformVendorEffectVibratorStep");
         try {
-            long vibratorOnResult = controller.on(effect, getVibration().id);
+            int stepId = conductor.nextVibratorCallbackStepId(getVibratorId());
+            long vibratorOnResult = controller.on(effect, getVibration().id, stepId);
             vibratorOnResult = Math.min(vibratorOnResult, VENDOR_EFFECT_MAX_DURATION_MS);
             handleVibratorOnResult(vibratorOnResult);
             getVibration().stats.reportPerformVendorEffect(vibratorOnResult);
diff --git a/services/core/java/com/android/server/vibrator/SetAmplitudeVibratorStep.java b/services/core/java/com/android/server/vibrator/SetAmplitudeVibratorStep.java
index 8478e7743183..26b9595e60cd 100644
--- a/services/core/java/com/android/server/vibrator/SetAmplitudeVibratorStep.java
+++ b/services/core/java/com/android/server/vibrator/SetAmplitudeVibratorStep.java
@@ -20,6 +20,7 @@ import android.annotation.NonNull;
 import android.os.SystemClock;
 import android.os.Trace;
 import android.os.VibrationEffect;
+import android.os.vibrator.Flags;
 import android.os.vibrator.StepSegment;
 import android.os.vibrator.VibrationEffectSegment;
 import android.util.Slog;
@@ -49,6 +50,10 @@ final class SetAmplitudeVibratorStep extends AbstractComposedVibratorStep {
 
     @Override
     public boolean acceptVibratorCompleteCallback(int vibratorId) {
+        if (Flags.fixVibrationThreadCallbackHandling()) {
+            // TODO: remove this method once flag removed.
+            return super.acceptVibratorCompleteCallback(vibratorId);
+        }
         // Ensure the super method is called and will reset the off timeout and boolean flag.
         // This is true if the vibrator was ON and this callback has the same vibratorId.
         if (!super.acceptVibratorCompleteCallback(vibratorId)) {
@@ -161,7 +166,8 @@ final class SetAmplitudeVibratorStep extends AbstractComposedVibratorStep {
                     "Turning on vibrator " + controller.getVibratorInfo().getId() + " for "
                             + duration + "ms");
         }
-        long vibratorOnResult = controller.on(duration, getVibration().id);
+        int stepId = conductor.nextVibratorCallbackStepId(getVibratorId());
+        long vibratorOnResult = controller.on(duration, getVibration().id, stepId);
         handleVibratorOnResult(vibratorOnResult);
         getVibration().stats.reportVibratorOn(vibratorOnResult);
         return vibratorOnResult;
diff --git a/services/core/java/com/android/server/vibrator/SingleVibrationSession.java b/services/core/java/com/android/server/vibrator/SingleVibrationSession.java
index 628221b09d77..309eb8c3b099 100644
--- a/services/core/java/com/android/server/vibrator/SingleVibrationSession.java
+++ b/services/core/java/com/android/server/vibrator/SingleVibrationSession.java
@@ -137,13 +137,13 @@ final class SingleVibrationSession implements VibrationSession, IBinder.DeathRec
     }
 
     @Override
-    public void notifyVibratorCallback(int vibratorId, long vibrationId) {
+    public void notifyVibratorCallback(int vibratorId, long vibrationId, long stepId) {
         if (vibrationId != mVibration.id) {
             return;
         }
         synchronized (mLock) {
             if (mConductor != null) {
-                mConductor.notifyVibratorComplete(vibratorId);
+                mConductor.notifyVibratorComplete(vibratorId, stepId);
             }
         }
     }
diff --git a/services/core/java/com/android/server/vibrator/VendorVibrationSession.java b/services/core/java/com/android/server/vibrator/VendorVibrationSession.java
index 64b52b175252..bda3d442956b 100644
--- a/services/core/java/com/android/server/vibrator/VendorVibrationSession.java
+++ b/services/core/java/com/android/server/vibrator/VendorVibrationSession.java
@@ -218,8 +218,8 @@ final class VendorVibrationSession extends IVibrationSession.Stub
     }
 
     @Override
-    public void notifyVibratorCallback(int vibratorId, long vibrationId) {
-        Slog.d(TAG, "Vibration callback received for vibration " + vibrationId
+    public void notifyVibratorCallback(int vibratorId, long vibrationId, long stepId) {
+        Slog.d(TAG, "Vibration callback received for vibration " + vibrationId + " step " + stepId
                 + " on vibrator " + vibratorId + ", ignoring...");
     }
 
diff --git a/services/core/java/com/android/server/vibrator/VibrationSession.java b/services/core/java/com/android/server/vibrator/VibrationSession.java
index ae95a70e2a4f..23715e392580 100644
--- a/services/core/java/com/android/server/vibrator/VibrationSession.java
+++ b/services/core/java/com/android/server/vibrator/VibrationSession.java
@@ -106,7 +106,7 @@ interface VibrationSession {
      * complete (e.g. on(), perform(), compose()). This does not mean the vibration is complete,
      * since its playback might have one or more interactions with the vibrator hardware.
      */
-    void notifyVibratorCallback(int vibratorId, long vibrationId);
+    void notifyVibratorCallback(int vibratorId, long vibrationId, long stepId);
 
     /**
      * Notify all synced vibrators have completed the last synchronized command during the playback
diff --git a/services/core/java/com/android/server/vibrator/VibrationStepConductor.java b/services/core/java/com/android/server/vibrator/VibrationStepConductor.java
index 1e20debe156d..36e13224a476 100644
--- a/services/core/java/com/android/server/vibrator/VibrationStepConductor.java
+++ b/services/core/java/com/android/server/vibrator/VibrationStepConductor.java
@@ -30,6 +30,7 @@ import android.os.vibrator.VibrationEffectSegment;
 import android.util.IntArray;
 import android.util.Slog;
 import android.util.SparseArray;
+import android.util.SparseIntArray;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.server.vibrator.VibrationSession.Status;
@@ -93,6 +94,8 @@ final class VibrationStepConductor {
     private final Object mLock = new Object();
     @GuardedBy("mLock")
     private final IntArray mSignalVibratorsComplete;
+    @GuardedBy("mLock")
+    private final SparseIntArray mSignalVibratorStepIds;
     @Nullable
     @GuardedBy("mLock")
     private Vibration.EndInfo mSignalCancel = null;
@@ -121,6 +124,8 @@ final class VibrationStepConductor {
         this.vibratorManagerHooks = vibratorManagerHooks;
         this.mSignalVibratorsComplete =
                 new IntArray(mDeviceAdapter.getAvailableVibratorIds().length);
+        this.mSignalVibratorStepIds =
+                new SparseIntArray(mDeviceAdapter.getAvailableVibratorIds().length);
     }
 
     @Nullable
@@ -418,7 +423,7 @@ final class VibrationStepConductor {
      * <p>This is a lightweight method intended to be called directly via native callbacks.
      * The state update is recorded for processing on the main execution thread (VibrationThread).
      */
-    public void notifyVibratorComplete(int vibratorId) {
+    public void notifyVibratorComplete(int vibratorId, long stepId) {
         // HAL callbacks may be triggered directly within HAL calls, so these notifications
         // could be on the VibrationThread as it calls the HAL, or some other executor later.
         // Therefore no thread assertion is made here.
@@ -428,6 +433,14 @@ final class VibrationStepConductor {
         }
 
         synchronized (mLock) {
+            if (Flags.fixVibrationThreadCallbackHandling()
+                    && mSignalVibratorStepIds.get(vibratorId) != stepId) {
+                if (DEBUG) {
+                    Slog.d(TAG, "Vibrator " + vibratorId + " callback for step=" + stepId
+                            + " ignored, current step=" + mSignalVibratorStepIds.get(vibratorId));
+                }
+                return;
+            }
             mSignalVibratorsComplete.add(vibratorId);
             mLock.notify();
         }
@@ -645,6 +658,26 @@ final class VibrationStepConductor {
         }
     }
 
+    /**
+     * Updates and returns the next step id value to be used in vibrator commands.
+     *
+     * <p>This new step id will be kept by this conductor to filter out old callbacks that might be
+     * triggered too late by the HAL, preventing them from affecting the ongoing vibration playback.
+     */
+    public int nextVibratorCallbackStepId(int vibratorId) {
+        if (!Flags.fixVibrationThreadCallbackHandling()) {
+            return 0;
+        }
+        if (Build.IS_DEBUGGABLE) {
+            expectIsVibrationThread(true);
+        }
+        synchronized (mLock) {
+            int stepId = mSignalVibratorStepIds.get(vibratorId) + 1;
+            mSignalVibratorStepIds.put(vibratorId, stepId);
+            return stepId;
+        }
+    }
+
     private static CombinedVibration.Sequential toSequential(CombinedVibration effect) {
         if (effect instanceof CombinedVibration.Sequential) {
             return (CombinedVibration.Sequential) effect;
diff --git a/services/core/java/com/android/server/vibrator/VibratorController.java b/services/core/java/com/android/server/vibrator/VibratorController.java
index acb31ceb4027..ab13b0e88d04 100644
--- a/services/core/java/com/android/server/vibrator/VibratorController.java
+++ b/services/core/java/com/android/server/vibrator/VibratorController.java
@@ -72,7 +72,7 @@ final class VibratorController {
     public interface OnVibrationCompleteListener {
 
         /** Callback triggered when an active vibration command is complete. */
-        void onComplete(int vibratorId, long vibrationId);
+        void onComplete(int vibratorId, long vibrationId, long stepId);
     }
 
     /** Representation of the vibrator state based on the interactions through this controller. */
@@ -285,11 +285,11 @@ final class VibratorController {
      * @return The positive duration of the vibration started, if successful, zero if the vibrator
      * do not support the input or a negative number if the operation failed.
      */
-    public long on(long milliseconds, long vibrationId) {
+    public long on(long milliseconds, long vibrationId, long stepId) {
         Trace.traceBegin(TRACE_TAG_VIBRATOR, "VibratorController#on");
         try {
             synchronized (mLock) {
-                long duration = mNativeWrapper.on(milliseconds, vibrationId);
+                long duration = mNativeWrapper.on(milliseconds, vibrationId, stepId);
                 if (duration > 0) {
                     mCurrentAmplitude = -1;
                     updateStateAndNotifyListenersLocked(VibratorState.VIBRATING);
@@ -310,7 +310,7 @@ final class VibratorController {
      * @return The positive duration of the vibration started, if successful, zero if the vibrator
      * do not support the input or a negative number if the operation failed.
      */
-    public long on(VibrationEffect.VendorEffect vendorEffect, long vibrationId) {
+    public long on(VibrationEffect.VendorEffect vendorEffect, long vibrationId, long stepId) {
         Trace.traceBegin(TRACE_TAG_VIBRATOR, "VibratorController#on (vendor)");
         synchronized (mLock) {
             Parcel vendorData = Parcel.obtain();
@@ -319,7 +319,7 @@ final class VibratorController {
                 vendorData.setDataPosition(0);
                 long duration = mNativeWrapper.performVendorEffect(vendorData,
                         vendorEffect.getEffectStrength(), vendorEffect.getScale(),
-                        vendorEffect.getAdaptiveScale(), vibrationId);
+                        vendorEffect.getAdaptiveScale(), vibrationId, stepId);
                 if (duration > 0) {
                     mCurrentAmplitude = -1;
                     updateStateAndNotifyListenersLocked(VibratorState.VIBRATING);
@@ -341,12 +341,12 @@ final class VibratorController {
      * @return The positive duration of the vibration started, if successful, zero if the vibrator
      * do not support the input or a negative number if the operation failed.
      */
-    public long on(PrebakedSegment prebaked, long vibrationId) {
+    public long on(PrebakedSegment prebaked, long vibrationId, long stepId) {
         Trace.traceBegin(TRACE_TAG_VIBRATOR, "VibratorController#on (Prebaked)");
         try {
             synchronized (mLock) {
                 long duration = mNativeWrapper.perform(prebaked.getEffectId(),
-                        prebaked.getEffectStrength(), vibrationId);
+                        prebaked.getEffectStrength(), vibrationId, stepId);
                 if (duration > 0) {
                     mCurrentAmplitude = -1;
                     updateStateAndNotifyListenersLocked(VibratorState.VIBRATING);
@@ -367,14 +367,14 @@ final class VibratorController {
      * @return The positive duration of the vibration started, if successful, zero if the vibrator
      * do not support the input or a negative number if the operation failed.
      */
-    public long on(PrimitiveSegment[] primitives, long vibrationId) {
+    public long on(PrimitiveSegment[] primitives, long vibrationId, long stepId) {
         Trace.traceBegin(TRACE_TAG_VIBRATOR, "VibratorController#on (Primitive)");
         try {
             if (!mVibratorInfo.hasCapability(IVibrator.CAP_COMPOSE_EFFECTS)) {
                 return 0;
             }
             synchronized (mLock) {
-                long duration = mNativeWrapper.compose(primitives, vibrationId);
+                long duration = mNativeWrapper.compose(primitives, vibrationId, stepId);
                 if (duration > 0) {
                     mCurrentAmplitude = -1;
                     updateStateAndNotifyListenersLocked(VibratorState.VIBRATING);
@@ -394,7 +394,7 @@ final class VibratorController {
      *
      * @return The duration of the effect playing, or 0 if unsupported.
      */
-    public long on(RampSegment[] primitives, long vibrationId) {
+    public long on(RampSegment[] primitives, long vibrationId, long stepId) {
         Trace.traceBegin(TRACE_TAG_VIBRATOR, "VibratorController#on (PWLE)");
         try {
             if (!mVibratorInfo.hasCapability(IVibrator.CAP_COMPOSE_PWLE_EFFECTS)) {
@@ -402,7 +402,8 @@ final class VibratorController {
             }
             synchronized (mLock) {
                 int braking = mVibratorInfo.getDefaultBraking();
-                long duration = mNativeWrapper.composePwle(primitives, braking, vibrationId);
+                long duration = mNativeWrapper.composePwle(
+                        primitives, braking, vibrationId, stepId);
                 if (duration > 0) {
                     mCurrentAmplitude = -1;
                     updateStateAndNotifyListenersLocked(VibratorState.VIBRATING);
@@ -422,14 +423,14 @@ final class VibratorController {
      *
      * @return The duration of the effect playing, or 0 if unsupported.
      */
-    public long on(PwlePoint[] pwlePoints, long vibrationId) {
+    public long on(PwlePoint[] pwlePoints, long vibrationId, long stepId) {
         Trace.traceBegin(TRACE_TAG_VIBRATOR, "VibratorController#on (PWLE v2)");
         try {
             if (!mVibratorInfo.hasCapability(IVibrator.CAP_COMPOSE_PWLE_EFFECTS_V2)) {
                 return 0;
             }
             synchronized (mLock) {
-                long duration = mNativeWrapper.composePwleV2(pwlePoints, vibrationId);
+                long duration = mNativeWrapper.composePwleV2(pwlePoints, vibrationId, stepId);
                 if (duration > 0) {
                     mCurrentAmplitude = -1;
                     updateStateAndNotifyListenersLocked(VibratorState.VIBRATING);
@@ -544,26 +545,27 @@ final class VibratorController {
 
         private static native boolean isAvailable(long nativePtr);
 
-        private static native long on(long nativePtr, long milliseconds, long vibrationId);
+        private static native long on(long nativePtr, long milliseconds, long vibrationId,
+                long stepId);
 
         private static native void off(long nativePtr);
 
         private static native void setAmplitude(long nativePtr, float amplitude);
 
         private static native long performEffect(long nativePtr, long effect, long strength,
-                long vibrationId);
+                long vibrationId, long stepId);
 
         private static native long performVendorEffect(long nativePtr, Parcel vendorData,
-                long strength, float scale, float adaptiveScale, long vibrationId);
+                long strength, float scale, float adaptiveScale, long vibrationId, long stepId);
 
         private static native long performComposedEffect(long nativePtr, PrimitiveSegment[] effect,
-                long vibrationId);
+                long vibrationId, long stepId);
 
         private static native long performPwleEffect(long nativePtr, RampSegment[] effect,
-                int braking, long vibrationId);
+                int braking, long vibrationId, long stepId);
 
         private static native long performPwleV2Effect(long nativePtr, PwlePoint[] effect,
-                long vibrationId);
+                long vibrationId, long stepId);
 
         private static native void setExternalControl(long nativePtr, boolean enabled);
 
@@ -595,8 +597,8 @@ final class VibratorController {
         }
 
         /** Turns vibrator on for given time. */
-        public long on(long milliseconds, long vibrationId) {
-            return on(mNativePtr, milliseconds, vibrationId);
+        public long on(long milliseconds, long vibrationId, long stepId) {
+            return on(mNativePtr, milliseconds, vibrationId, stepId);
         }
 
         /** Turns vibrator off. */
@@ -610,30 +612,31 @@ final class VibratorController {
         }
 
         /** Turns vibrator on to perform one of the supported effects. */
-        public long perform(long effect, long strength, long vibrationId) {
-            return performEffect(mNativePtr, effect, strength, vibrationId);
+        public long perform(long effect, long strength, long vibrationId, long stepId) {
+            return performEffect(mNativePtr, effect, strength, vibrationId, stepId);
         }
 
         /** Turns vibrator on to perform a vendor-specific effect. */
         public long performVendorEffect(Parcel vendorData, long strength, float scale,
-                float adaptiveScale, long vibrationId) {
+                float adaptiveScale, long vibrationId, long stepId) {
             return performVendorEffect(mNativePtr, vendorData, strength, scale, adaptiveScale,
-                    vibrationId);
+                    vibrationId, stepId);
         }
 
         /** Turns vibrator on to perform effect composed of give primitives effect. */
-        public long compose(PrimitiveSegment[] primitives, long vibrationId) {
-            return performComposedEffect(mNativePtr, primitives, vibrationId);
+        public long compose(PrimitiveSegment[] primitives, long vibrationId, long stepId) {
+            return performComposedEffect(mNativePtr, primitives, vibrationId, stepId);
         }
 
         /** Turns vibrator on to perform PWLE effect composed of given primitives. */
-        public long composePwle(RampSegment[] primitives, int braking, long vibrationId) {
-            return performPwleEffect(mNativePtr, primitives, braking, vibrationId);
+        public long composePwle(RampSegment[] primitives, int braking, long vibrationId,
+                long stepId) {
+            return performPwleEffect(mNativePtr, primitives, braking, vibrationId, stepId);
         }
 
         /** Turns vibrator on to perform PWLE effect composed of given points. */
-        public long composePwleV2(PwlePoint[] pwlePoints, long vibrationId) {
-            return performPwleV2Effect(mNativePtr, pwlePoints, vibrationId);
+        public long composePwleV2(PwlePoint[] pwlePoints, long vibrationId, long stepId) {
+            return performPwleV2Effect(mNativePtr, pwlePoints, vibrationId, stepId);
         }
 
         /** Enabled the device vibrator to be controlled by another service. */
diff --git a/services/core/java/com/android/server/vibrator/VibratorManagerService.java b/services/core/java/com/android/server/vibrator/VibratorManagerService.java
index 75b1b202bcfd..3f5fc338ee3b 100644
--- a/services/core/java/com/android/server/vibrator/VibratorManagerService.java
+++ b/services/core/java/com/android/server/vibrator/VibratorManagerService.java
@@ -1293,14 +1293,14 @@ public class VibratorManagerService extends IVibratorManagerService.Stub {
         }
     }
 
-    private void onVibrationComplete(int vibratorId, long vibrationId) {
+    private void onVibrationComplete(int vibratorId, long vibrationId, long stepId) {
         synchronized (mLock) {
             if (mCurrentSession != null) {
                 if (DEBUG) {
-                    Slog.d(TAG, "Vibration " + vibrationId + " on vibrator " + vibratorId
-                            + " complete, notifying thread");
+                    Slog.d(TAG, "Vibration " + vibrationId + " step " + stepId
+                            + " on vibrator " + vibratorId + " complete, notifying thread");
                 }
-                mCurrentSession.notifyVibratorCallback(vibratorId, vibrationId);
+                mCurrentSession.notifyVibratorCallback(vibratorId, vibrationId, stepId);
             }
         }
     }
@@ -2100,10 +2100,10 @@ public class VibratorManagerService extends IVibratorManagerService.Stub {
         }
 
         @Override
-        public void onComplete(int vibratorId, long vibrationId) {
+        public void onComplete(int vibratorId, long vibrationId, long stepId) {
             VibratorManagerService service = mServiceRef.get();
             if (service != null) {
-                service.onVibrationComplete(vibratorId, vibrationId);
+                service.onVibrationComplete(vibratorId, vibrationId, stepId);
             }
         }
     }
diff --git a/services/core/jni/Android.bp b/services/core/jni/Android.bp
index d26539c377a9..95776088aad8 100644
--- a/services/core/jni/Android.bp
+++ b/services/core/jni/Android.bp
@@ -204,6 +204,7 @@ cc_defaults {
         "android.frameworks.sensorservice-V1-ndk",
         "android.frameworks.stats@1.0",
         "android.frameworks.stats-V2-ndk",
+        "android.os.vibrator.flags-aconfig-cc",
         "android.system.suspend.control-V1-cpp",
         "android.system.suspend.control.internal-cpp",
         "android.system.suspend-V1-ndk",
diff --git a/services/core/jni/com_android_server_vibrator_VibratorController.cpp b/services/core/jni/com_android_server_vibrator_VibratorController.cpp
index abd4cd25cf68..534dbb1f6cf1 100644
--- a/services/core/jni/com_android_server_vibrator_VibratorController.cpp
+++ b/services/core/jni/com_android_server_vibrator_VibratorController.cpp
@@ -21,6 +21,7 @@
 #include <android/binder_parcel_jni.h>
 #include <android/hardware/vibrator/1.3/IVibrator.h>
 #include <android/persistable_bundle_aidl.h>
+#include <android_os_vibrator.h>
 #include <nativehelper/JNIHelp.h>
 #include <utils/Log.h>
 #include <utils/misc.h>
@@ -143,21 +144,23 @@ public:
         return mHal->doWithRetry(fn, functionName);
     }
 
-    std::function<void()> createCallback(jlong vibrationId) {
+    std::function<void()> createCallback(jlong vibrationId, jlong stepId) {
         auto callbackId = ++mCallbackId;
-        return [vibrationId, callbackId, this]() {
+        return [vibrationId, stepId, callbackId, this]() {
             auto currentCallbackId = mCallbackId.load();
-            if (currentCallbackId != callbackId) {
-                // This callback is from an older HAL call that is no longer relevant to the service
+            if (!android_os_vibrator_fix_vibration_thread_callback_handling() &&
+                currentCallbackId != callbackId) {
+                // This callback is from an older HAL call that is no longer relevant
                 return;
             }
             auto jniEnv = GetOrAttachJNIEnvironment(sJvm);
-            jniEnv->CallVoidMethod(mCallbackListener, sMethodIdOnComplete, mVibratorId,
-                                   vibrationId);
+            jniEnv->CallVoidMethod(mCallbackListener, sMethodIdOnComplete, mVibratorId, vibrationId,
+                                   stepId);
         };
     }
 
     void disableOldCallbacks() {
+        // TODO remove this once android_os_vibrator_fix_vibration_thread_callback_handling removed
         mCallbackId++;
     }
 
@@ -165,6 +168,7 @@ private:
     const std::shared_ptr<vibrator::HalController> mHal;
     const int32_t mVibratorId;
     const jobject mCallbackListener;
+    // TODO remove this once android_os_vibrator_fix_vibration_thread_callback_handling removed
     std::atomic<int64_t> mCallbackId;
 };
 
@@ -273,13 +277,13 @@ static jboolean vibratorIsAvailable(JNIEnv* env, jclass /* clazz */, jlong ptr)
 }
 
 static jlong vibratorOn(JNIEnv* env, jclass /* clazz */, jlong ptr, jlong timeoutMs,
-                        jlong vibrationId) {
+                        jlong vibrationId, jlong stepId) {
     VibratorControllerWrapper* wrapper = reinterpret_cast<VibratorControllerWrapper*>(ptr);
     if (wrapper == nullptr) {
         ALOGE("vibratorOn failed because native wrapper was not initialized");
         return -1;
     }
-    auto callback = wrapper->createCallback(vibrationId);
+    auto callback = wrapper->createCallback(vibrationId, stepId);
     auto onFn = [timeoutMs, &callback](vibrator::HalWrapper* hal) {
         return hal->on(std::chrono::milliseconds(timeoutMs), callback);
     };
@@ -324,7 +328,7 @@ static void vibratorSetExternalControl(JNIEnv* env, jclass /* clazz */, jlong pt
 }
 
 static jlong vibratorPerformEffect(JNIEnv* env, jclass /* clazz */, jlong ptr, jlong effect,
-                                   jlong strength, jlong vibrationId) {
+                                   jlong strength, jlong vibrationId, jlong stepId) {
     VibratorControllerWrapper* wrapper = reinterpret_cast<VibratorControllerWrapper*>(ptr);
     if (wrapper == nullptr) {
         ALOGE("vibratorPerformEffect failed because native wrapper was not initialized");
@@ -332,7 +336,7 @@ static jlong vibratorPerformEffect(JNIEnv* env, jclass /* clazz */, jlong ptr, j
     }
     Aidl::Effect effectType = static_cast<Aidl::Effect>(effect);
     Aidl::EffectStrength effectStrength = static_cast<Aidl::EffectStrength>(strength);
-    auto callback = wrapper->createCallback(vibrationId);
+    auto callback = wrapper->createCallback(vibrationId, stepId);
     auto performEffectFn = [effectType, effectStrength, &callback](vibrator::HalWrapper* hal) {
         return hal->performEffect(effectType, effectStrength, callback);
     };
@@ -342,7 +346,7 @@ static jlong vibratorPerformEffect(JNIEnv* env, jclass /* clazz */, jlong ptr, j
 
 static jlong vibratorPerformVendorEffect(JNIEnv* env, jclass /* clazz */, jlong ptr,
                                          jobject vendorData, jlong strength, jfloat scale,
-                                         jfloat adaptiveScale, jlong vibrationId) {
+                                         jfloat adaptiveScale, jlong vibrationId, jlong stepId) {
     VibratorControllerWrapper* wrapper = reinterpret_cast<VibratorControllerWrapper*>(ptr);
     if (wrapper == nullptr) {
         ALOGE("vibratorPerformVendorEffect failed because native wrapper was not initialized");
@@ -350,7 +354,7 @@ static jlong vibratorPerformVendorEffect(JNIEnv* env, jclass /* clazz */, jlong
     }
     Aidl::VendorEffect effect =
             vendorEffectFromJavaParcel(env, vendorData, strength, scale, adaptiveScale);
-    auto callback = wrapper->createCallback(vibrationId);
+    auto callback = wrapper->createCallback(vibrationId, stepId);
     auto performVendorEffectFn = [&effect, &callback](vibrator::HalWrapper* hal) {
         return hal->performVendorEffect(effect, callback);
     };
@@ -359,7 +363,8 @@ static jlong vibratorPerformVendorEffect(JNIEnv* env, jclass /* clazz */, jlong
 }
 
 static jlong vibratorPerformComposedEffect(JNIEnv* env, jclass /* clazz */, jlong ptr,
-                                           jobjectArray composition, jlong vibrationId) {
+                                           jobjectArray composition, jlong vibrationId,
+                                           jlong stepId) {
     VibratorControllerWrapper* wrapper = reinterpret_cast<VibratorControllerWrapper*>(ptr);
     if (wrapper == nullptr) {
         ALOGE("vibratorPerformComposedEffect failed because native wrapper was not initialized");
@@ -371,7 +376,7 @@ static jlong vibratorPerformComposedEffect(JNIEnv* env, jclass /* clazz */, jlon
         jobject element = env->GetObjectArrayElement(composition, i);
         effects.push_back(effectFromJavaPrimitive(env, element));
     }
-    auto callback = wrapper->createCallback(vibrationId);
+    auto callback = wrapper->createCallback(vibrationId, stepId);
     auto performComposedEffectFn = [&effects, &callback](vibrator::HalWrapper* hal) {
         return hal->performComposedEffect(effects, callback);
     };
@@ -381,7 +386,8 @@ static jlong vibratorPerformComposedEffect(JNIEnv* env, jclass /* clazz */, jlon
 }
 
 static jlong vibratorPerformPwleEffect(JNIEnv* env, jclass /* clazz */, jlong ptr,
-                                       jobjectArray waveform, jint brakingId, jlong vibrationId) {
+                                       jobjectArray waveform, jint brakingId, jlong vibrationId,
+                                       jlong stepId) {
     VibratorControllerWrapper* wrapper = reinterpret_cast<VibratorControllerWrapper*>(ptr);
     if (wrapper == nullptr) {
         ALOGE("vibratorPerformPwleEffect failed because native wrapper was not initialized");
@@ -406,7 +412,7 @@ static jlong vibratorPerformPwleEffect(JNIEnv* env, jclass /* clazz */, jlong pt
         }
     }
 
-    auto callback = wrapper->createCallback(vibrationId);
+    auto callback = wrapper->createCallback(vibrationId, stepId);
     auto performPwleEffectFn = [&primitives, &callback](vibrator::HalWrapper* hal) {
         return hal->performPwleEffect(primitives, callback);
     };
@@ -415,7 +421,7 @@ static jlong vibratorPerformPwleEffect(JNIEnv* env, jclass /* clazz */, jlong pt
 }
 
 static jlong vibratorPerformPwleV2Effect(JNIEnv* env, jclass /* clazz */, jlong ptr,
-                                         jobjectArray waveform, jlong vibrationId) {
+                                         jobjectArray waveform, jlong vibrationId, jlong stepId) {
     VibratorControllerWrapper* wrapper = reinterpret_cast<VibratorControllerWrapper*>(ptr);
     if (wrapper == nullptr) {
         ALOGE("vibratorPerformPwleV2Effect failed because native wrapper was not initialized");
@@ -431,7 +437,7 @@ static jlong vibratorPerformPwleV2Effect(JNIEnv* env, jclass /* clazz */, jlong
     }
     composite.pwlePrimitives = primitives;
 
-    auto callback = wrapper->createCallback(vibrationId);
+    auto callback = wrapper->createCallback(vibrationId, stepId);
     auto composePwleV2Fn = [&composite, &callback](vibrator::HalWrapper* hal) {
         return hal->composePwleV2(composite, callback);
     };
@@ -610,16 +616,16 @@ static const JNINativeMethod method_table[] = {
          (void*)vibratorNativeInit},
         {"getNativeFinalizer", "()J", (void*)vibratorGetNativeFinalizer},
         {"isAvailable", "(J)Z", (void*)vibratorIsAvailable},
-        {"on", "(JJJ)J", (void*)vibratorOn},
+        {"on", "(JJJJ)J", (void*)vibratorOn},
         {"off", "(J)V", (void*)vibratorOff},
         {"setAmplitude", "(JF)V", (void*)vibratorSetAmplitude},
-        {"performEffect", "(JJJJ)J", (void*)vibratorPerformEffect},
-        {"performVendorEffect", "(JLandroid/os/Parcel;JFFJ)J", (void*)vibratorPerformVendorEffect},
-        {"performComposedEffect", "(J[Landroid/os/vibrator/PrimitiveSegment;J)J",
+        {"performEffect", "(JJJJJ)J", (void*)vibratorPerformEffect},
+        {"performVendorEffect", "(JLandroid/os/Parcel;JFFJJ)J", (void*)vibratorPerformVendorEffect},
+        {"performComposedEffect", "(J[Landroid/os/vibrator/PrimitiveSegment;JJ)J",
          (void*)vibratorPerformComposedEffect},
-        {"performPwleEffect", "(J[Landroid/os/vibrator/RampSegment;IJ)J",
+        {"performPwleEffect", "(J[Landroid/os/vibrator/RampSegment;IJJ)J",
          (void*)vibratorPerformPwleEffect},
-        {"performPwleV2Effect", "(J[Landroid/os/vibrator/PwlePoint;J)J",
+        {"performPwleV2Effect", "(J[Landroid/os/vibrator/PwlePoint;JJ)J",
          (void*)vibratorPerformPwleV2Effect},
         {"setExternalControl", "(JZ)V", (void*)vibratorSetExternalControl},
         {"alwaysOnEnable", "(JJJJ)V", (void*)vibratorAlwaysOnEnable},
@@ -632,7 +638,7 @@ int register_android_server_vibrator_VibratorController(JavaVM* jvm, JNIEnv* env
     auto listenerClassName =
             "com/android/server/vibrator/VibratorController$OnVibrationCompleteListener";
     jclass listenerClass = FindClassOrDie(env, listenerClassName);
-    sMethodIdOnComplete = GetMethodIDOrDie(env, listenerClass, "onComplete", "(IJ)V");
+    sMethodIdOnComplete = GetMethodIDOrDie(env, listenerClass, "onComplete", "(IJJ)V");
 
     jclass primitiveClass = FindClassOrDie(env, "android/os/vibrator/PrimitiveSegment");
     sPrimitiveClassInfo.id = GetFieldIDOrDie(env, primitiveClass, "mPrimitiveId", "I");
diff --git a/services/tests/vibrator/src/com/android/server/vibrator/DeviceAdapterTest.java b/services/tests/vibrator/src/com/android/server/vibrator/DeviceAdapterTest.java
index d5548a4f375e..f091a65cfe46 100644
--- a/services/tests/vibrator/src/com/android/server/vibrator/DeviceAdapterTest.java
+++ b/services/tests/vibrator/src/com/android/server/vibrator/DeviceAdapterTest.java
@@ -497,19 +497,19 @@ public class DeviceAdapterTest {
 
     private VibratorController createEmptyVibratorController(int vibratorId) {
         return new FakeVibratorControllerProvider(mTestLooper.getLooper())
-                .newVibratorController(vibratorId, (id, vibrationId)  -> {});
+                .newVibratorController(vibratorId, (id, vibrationId, stepId)  -> {});
     }
 
     private VibratorController createBasicVibratorController(int vibratorId) {
         FakeVibratorControllerProvider provider = createVibratorProviderWithEffects(
                 IVibrator.CAP_COMPOSE_EFFECTS);
-        return provider.newVibratorController(vibratorId, (id, vibrationId)  -> {});
+        return provider.newVibratorController(vibratorId, (id, vibrationId, stepId)  -> {});
     }
 
     private VibratorController createPwleWithoutFrequenciesVibratorController(int vibratorId) {
         FakeVibratorControllerProvider provider = createVibratorProviderWithEffects(
                 IVibrator.CAP_COMPOSE_EFFECTS, IVibrator.CAP_COMPOSE_PWLE_EFFECTS);
-        return provider.newVibratorController(vibratorId, (id, vibrationId)  -> {});
+        return provider.newVibratorController(vibratorId, (id, vibrationId, stepId)  -> {});
     }
 
     private VibratorController createPwleVibratorController(int vibratorId) {
@@ -519,7 +519,7 @@ public class DeviceAdapterTest {
         provider.setMinFrequency(TEST_MIN_FREQUENCY);
         provider.setFrequencyResolution(TEST_FREQUENCY_RESOLUTION);
         provider.setMaxAmplitudes(TEST_AMPLITUDE_MAP);
-        return provider.newVibratorController(vibratorId, (id, vibrationId)  -> {});
+        return provider.newVibratorController(vibratorId, (id, vibrationId, stepId)  -> {});
     }
 
     private VibratorController createPwleV2VibratorController(int vibratorId) {
@@ -538,7 +538,7 @@ public class DeviceAdapterTest {
         provider.setMinEnvelopeEffectControlPointDurationMillis(
                 TEST_MIN_ENVELOPE_EFFECT_CONTROL_POINT_DURATION_MILLIS);
 
-        return provider.newVibratorController(vibratorId, (id, vibrationId)  -> {});
+        return provider.newVibratorController(vibratorId, (id, vibrationId, stepId)  -> {});
     }
 
     private FakeVibratorControllerProvider createVibratorProviderWithEffects(int... capabilities) {
diff --git a/services/tests/vibrator/src/com/android/server/vibrator/VibrationThreadTest.java b/services/tests/vibrator/src/com/android/server/vibrator/VibrationThreadTest.java
index 5f2af0a085c3..42279e40fa33 100644
--- a/services/tests/vibrator/src/com/android/server/vibrator/VibrationThreadTest.java
+++ b/services/tests/vibrator/src/com/android/server/vibrator/VibrationThreadTest.java
@@ -17,6 +17,9 @@
 package com.android.server.vibrator;
 
 import static android.os.VibrationAttributes.USAGE_RINGTONE;
+import static android.os.VibrationEffect.Composition.PRIMITIVE_CLICK;
+import static android.os.VibrationEffect.Composition.PRIMITIVE_SPIN;
+import static android.os.VibrationEffect.Composition.PRIMITIVE_TICK;
 import static android.os.VibrationEffect.VibrationParameter.targetAmplitude;
 import static android.os.VibrationEffect.VibrationParameter.targetFrequency;
 
@@ -66,7 +69,6 @@ import android.os.vibrator.VibrationConfig;
 import android.os.vibrator.VibrationEffectSegment;
 import android.platform.test.annotations.DisableFlags;
 import android.platform.test.annotations.EnableFlags;
-import android.platform.test.annotations.RequiresFlagsEnabled;
 import android.platform.test.flag.junit.CheckFlagsRule;
 import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.platform.test.flag.junit.SetFlagsRule;
@@ -153,9 +155,10 @@ public class VibrationThreadTest {
         when(mPackageManagerInternalMock.getSystemUiServiceComponent())
                 .thenReturn(new ComponentName("", ""));
         doAnswer(answer -> {
-            mVibrationConductor.notifyVibratorComplete(answer.getArgument(0));
+            mVibrationConductor.notifyVibratorComplete(
+                    answer.getArgument(0), answer.getArgument(2));
             return null;
-        }).when(mControllerCallbacks).onComplete(anyInt(), anyLong());
+        }).when(mControllerCallbacks).onComplete(anyInt(), anyLong(), anyLong());
 
         LocalServices.removeServiceForTest(PackageManagerInternal.class);
         LocalServices.addService(PackageManagerInternal.class, mPackageManagerInternalMock);
@@ -190,7 +193,7 @@ public class VibrationThreadTest {
         HalVibration vibration = startThreadAndDispatcher(effect);
         waitForCompletion();
 
-        verify(mControllerCallbacks, never()).onComplete(anyInt(), eq(vibration.id));
+        verify(mControllerCallbacks, never()).onComplete(anyInt(), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.IGNORED_UNSUPPORTED);
     }
 
@@ -203,7 +206,7 @@ public class VibrationThreadTest {
         HalVibration vibration = startThreadAndDispatcher(effect);
         waitForCompletion();
 
-        verify(mControllerCallbacks, never()).onComplete(anyInt(), eq(vibration.id));
+        verify(mControllerCallbacks, never()).onComplete(anyInt(), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.IGNORED_UNSUPPORTED);
     }
 
@@ -217,7 +220,7 @@ public class VibrationThreadTest {
 
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(10L));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
-        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id));
+        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.FINISHED);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
 
@@ -234,7 +237,7 @@ public class VibrationThreadTest {
 
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(10L));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
-        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id));
+        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.FINISHED);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
 
@@ -254,7 +257,7 @@ public class VibrationThreadTest {
 
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(15L));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
-        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id));
+        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.FINISHED);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
 
@@ -265,7 +268,7 @@ public class VibrationThreadTest {
     }
 
     @Test
-    @RequiresFlagsEnabled(Flags.FLAG_ADAPTIVE_HAPTICS_ENABLED)
+    @EnableFlags(Flags.FLAG_ADAPTIVE_HAPTICS_ENABLED)
     public void vibrate_singleWaveformWithAdaptiveHapticsScaling_scalesAmplitudesProperly() {
         // No user settings scale.
         setUserSetting(Settings.System.RING_VIBRATION_INTENSITY,
@@ -291,7 +294,7 @@ public class VibrationThreadTest {
     }
 
     @Test
-    @RequiresFlagsEnabled(Flags.FLAG_ADAPTIVE_HAPTICS_ENABLED)
+    @EnableFlags(Flags.FLAG_ADAPTIVE_HAPTICS_ENABLED)
     public void vibrate_withVibrationParamsRequestStalling_timeoutRequestAndApplyNoScaling() {
         // No user settings scale.
         setUserSetting(Settings.System.RING_VIBRATION_INTENSITY,
@@ -412,6 +415,7 @@ public class VibrationThreadTest {
                 .isEqualTo(expectedOneShots(200L, 50L));
     }
 
+    @EnableFlags(Flags.FLAG_FIX_VIBRATION_THREAD_CALLBACK_HANDLING)
     @LargeTest
     @Test
     public void vibrate_singleVibratorRepeatingPatternWithZeroDurationSteps_repeatsEffectCorrectly()
@@ -443,6 +447,30 @@ public class VibrationThreadTest {
                 .isEqualTo(expectedOneShots(200L, 150L, 300L, 150L, 300L));
     }
 
+    @EnableFlags(Flags.FLAG_FIX_VIBRATION_THREAD_CALLBACK_HANDLING)
+    @Test
+    public void vibrate_singleVibratorPatternWithCallbackDelay_oldCallbacksIgnored() {
+        FakeVibratorControllerProvider fakeVibrator = mVibratorProviders.get(VIBRATOR_ID);
+        fakeVibrator.setCompletionCallbackDelay(100); // 100ms delay to notify service.
+        fakeVibrator.setCapabilities(IVibrator.CAP_AMPLITUDE_CONTROL);
+
+        VibrationEffect effect = VibrationEffect.createWaveform(
+                /* timings= */ new long[]{0, 200, 50, 400}, /* repeat= */ -1);
+        HalVibration vibration = startThreadAndDispatcher(effect);
+        waitForCompletion(800 + TEST_TIMEOUT_MILLIS); // 200 + 50 + 400 + 100ms delay
+
+        verifyCallbacksTriggered(vibration, Status.FINISHED);
+        assertThat(mControllers.get(VIBRATOR_ID).isVibrating()).isFalse();
+
+        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id), eq(1L));
+        // Step id = 2 skipped by the 50ms OFF step after the 200ms ON step.
+        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id), eq(3L));
+
+        // First callback ignored, did not cause the vibrator to turn back on during the 400ms step.
+        assertThat(fakeVibrator.getEffectSegments(vibration.id))
+                .isEqualTo(expectedOneShots(200L, 400L));
+    }
+
     @Test
     public void vibrate_singleVibratorRepeatingPwle_generatesLargestPwles() throws Exception {
         FakeVibratorControllerProvider fakeVibrator = mVibratorProviders.get(VIBRATOR_ID);
@@ -479,12 +507,12 @@ public class VibrationThreadTest {
             throws Exception {
         FakeVibratorControllerProvider fakeVibrator = mVibratorProviders.get(VIBRATOR_ID);
         fakeVibrator.setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS);
-        fakeVibrator.setSupportedPrimitives(VibrationEffect.Composition.PRIMITIVE_CLICK);
+        fakeVibrator.setSupportedPrimitives(PRIMITIVE_CLICK);
         fakeVibrator.setCompositionSizeMax(10);
 
         VibrationEffect effect = VibrationEffect.startComposition()
                 // Very long delay so thread will be cancelled after first PWLE is triggered.
-                .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 1f, 100)
+                .addPrimitive(PRIMITIVE_CLICK, 1f, 100)
                 .compose();
         VibrationEffect repeatingEffect = VibrationEffect.startComposition()
                 .repeatEffectIndefinitely(effect)
@@ -561,13 +589,12 @@ public class VibrationThreadTest {
     public void vibrate_singleVibratorPredefinedCancel_cancelsVibrationImmediately()
             throws Exception {
         mVibratorProviders.get(VIBRATOR_ID).setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS);
-        mVibratorProviders.get(VIBRATOR_ID).setSupportedPrimitives(
-                VibrationEffect.Composition.PRIMITIVE_CLICK);
+        mVibratorProviders.get(VIBRATOR_ID).setSupportedPrimitives(PRIMITIVE_CLICK);
 
         VibrationEffect effect = VibrationEffect.startComposition()
-                .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 1f, 100)
-                .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 1f, 100)
-                .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 1f, 100)
+                .addPrimitive(PRIMITIVE_CLICK, 1f, 100)
+                .addPrimitive(PRIMITIVE_CLICK, 1f, 100)
+                .addPrimitive(PRIMITIVE_CLICK, 1f, 100)
                 .compose();
         HalVibration vibration = startThreadAndDispatcher(effect);
 
@@ -591,7 +618,7 @@ public class VibrationThreadTest {
     }
 
     @Test
-    @RequiresFlagsEnabled(android.os.vibrator.Flags.FLAG_VENDOR_VIBRATION_EFFECTS)
+    @EnableFlags(Flags.FLAG_VENDOR_VIBRATION_EFFECTS)
     public void vibrate_singleVibratorVendorEffectCancel_cancelsVibrationImmediately()
             throws Exception {
         mVibratorProviders.get(VIBRATOR_ID).setCapabilities(IVibrator.CAP_PERFORM_VENDOR_EFFECTS);
@@ -657,7 +684,7 @@ public class VibrationThreadTest {
 
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(20L));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
-        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id));
+        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.FINISHED);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
 
@@ -678,7 +705,7 @@ public class VibrationThreadTest {
 
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(10L));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
-        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id));
+        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.FINISHED);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
 
@@ -695,13 +722,14 @@ public class VibrationThreadTest {
 
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(0L));
         verify(mManagerHooks, never()).noteVibratorOff(eq(UID));
-        verify(mControllerCallbacks, never()).onComplete(eq(VIBRATOR_ID), eq(vibration.id));
+        verify(mControllerCallbacks, never())
+                .onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.IGNORED_UNSUPPORTED);
         assertTrue(mVibratorProviders.get(VIBRATOR_ID).getEffectSegments(vibration.id).isEmpty());
     }
 
     @Test
-    @RequiresFlagsEnabled(android.os.vibrator.Flags.FLAG_VENDOR_VIBRATION_EFFECTS)
+    @EnableFlags(Flags.FLAG_VENDOR_VIBRATION_EFFECTS)
     public void vibrate_singleVibratorVendorEffect_runsVibration() {
         mVibratorProviders.get(1).setCapabilities(IVibrator.CAP_PERFORM_VENDOR_EFFECTS);
 
@@ -712,7 +740,7 @@ public class VibrationThreadTest {
         verify(mManagerHooks).noteVibratorOn(eq(UID),
                 eq(PerformVendorEffectVibratorStep.VENDOR_EFFECT_MAX_DURATION_MS));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
-        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id));
+        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.FINISHED);
         assertThat(mControllers.get(VIBRATOR_ID).isVibrating()).isFalse();
 
@@ -725,24 +753,23 @@ public class VibrationThreadTest {
     public void vibrate_singleVibratorComposed_runsVibration() {
         FakeVibratorControllerProvider fakeVibrator = mVibratorProviders.get(VIBRATOR_ID);
         fakeVibrator.setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS);
-        fakeVibrator.setSupportedPrimitives(VibrationEffect.Composition.PRIMITIVE_CLICK,
-                VibrationEffect.Composition.PRIMITIVE_TICK);
+        fakeVibrator.setSupportedPrimitives(PRIMITIVE_CLICK, PRIMITIVE_TICK);
 
         VibrationEffect effect = VibrationEffect.startComposition()
-                .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 1f)
-                .addPrimitive(VibrationEffect.Composition.PRIMITIVE_TICK, 0.5f)
+                .addPrimitive(PRIMITIVE_CLICK, 1f)
+                .addPrimitive(PRIMITIVE_TICK, 0.5f)
                 .compose();
         HalVibration vibration = startThreadAndDispatcher(effect);
         waitForCompletion();
 
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(40L));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
-        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id));
+        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.FINISHED);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
         assertEquals(Arrays.asList(
-                expectedPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 1, 0),
-                expectedPrimitive(VibrationEffect.Composition.PRIMITIVE_TICK, 0.5f, 0)),
+                expectedPrimitive(PRIMITIVE_CLICK, 1, 0),
+                expectedPrimitive(PRIMITIVE_TICK, 0.5f, 0)),
                 fakeVibrator.getEffectSegments(vibration.id));
     }
 
@@ -750,14 +777,15 @@ public class VibrationThreadTest {
     @DisableFlags(Flags.FLAG_PRIMITIVE_COMPOSITION_ABSOLUTE_DELAY)
     public void vibrate_singleVibratorComposedAndNoCapability_triggersHalAndReturnsUnsupported() {
         VibrationEffect effect = VibrationEffect.startComposition()
-                .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 1f)
+                .addPrimitive(PRIMITIVE_CLICK, 1f)
                 .compose();
         HalVibration vibration = startThreadAndDispatcher(effect);
         waitForCompletion();
 
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(0L));
         verify(mManagerHooks, never()).noteVibratorOff(eq(UID));
-        verify(mControllerCallbacks, never()).onComplete(eq(VIBRATOR_ID), eq(vibration.id));
+        verify(mControllerCallbacks, never())
+                .onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.IGNORED_UNSUPPORTED);
         assertTrue(mVibratorProviders.get(VIBRATOR_ID).getEffectSegments(vibration.id).isEmpty());
     }
@@ -766,14 +794,15 @@ public class VibrationThreadTest {
     @EnableFlags(Flags.FLAG_PRIMITIVE_COMPOSITION_ABSOLUTE_DELAY)
     public void vibrate_singleVibratorComposedAndNoCapability_ignoresVibration() {
         VibrationEffect effect = VibrationEffect.startComposition()
-                .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 1f)
+                .addPrimitive(PRIMITIVE_CLICK, 1f)
                 .compose();
         HalVibration vibration = startThreadAndDispatcher(effect);
         waitForCompletion();
 
         verify(mManagerHooks, never()).noteVibratorOn(eq(UID), anyLong());
         verify(mManagerHooks, never()).noteVibratorOff(eq(UID));
-        verify(mControllerCallbacks, never()).onComplete(eq(VIBRATOR_ID), eq(vibration.id));
+        verify(mControllerCallbacks, never())
+                .onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.IGNORED_UNSUPPORTED);
         assertTrue(mVibratorProviders.get(VIBRATOR_ID).getEffectSegments(vibration.id).isEmpty());
     }
@@ -782,34 +811,30 @@ public class VibrationThreadTest {
     public void vibrate_singleVibratorLargeComposition_splitsVibratorComposeCalls() {
         FakeVibratorControllerProvider fakeVibrator = mVibratorProviders.get(VIBRATOR_ID);
         fakeVibrator.setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS);
-        fakeVibrator.setSupportedPrimitives(
-                VibrationEffect.Composition.PRIMITIVE_CLICK,
-                VibrationEffect.Composition.PRIMITIVE_TICK,
-                VibrationEffect.Composition.PRIMITIVE_SPIN);
+        fakeVibrator.setSupportedPrimitives(PRIMITIVE_CLICK, PRIMITIVE_TICK, PRIMITIVE_SPIN);
         fakeVibrator.setCompositionSizeMax(2);
 
         VibrationEffect effect = VibrationEffect.startComposition()
-                .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 1f)
-                .addPrimitive(VibrationEffect.Composition.PRIMITIVE_TICK, 0.5f)
-                .addPrimitive(VibrationEffect.Composition.PRIMITIVE_SPIN, 0.8f)
+                .addPrimitive(PRIMITIVE_CLICK, 1f)
+                .addPrimitive(PRIMITIVE_TICK, 0.5f)
+                .addPrimitive(PRIMITIVE_SPIN, 0.8f)
                 .compose();
         HalVibration vibration = startThreadAndDispatcher(effect);
         waitForCompletion();
 
         verifyCallbacksTriggered(vibration, Status.FINISHED);
         // Vibrator compose called twice.
-        verify(mControllerCallbacks, times(2)).onComplete(eq(VIBRATOR_ID), eq(vibration.id));
+        verify(mControllerCallbacks, times(2))
+                .onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         assertEquals(3, fakeVibrator.getEffectSegments(vibration.id).size());
     }
 
     @Test
-    @DisableFlags(android.os.vibrator.Flags.FLAG_NORMALIZED_PWLE_EFFECTS)
+    @DisableFlags(Flags.FLAG_NORMALIZED_PWLE_EFFECTS)
     public void vibrate_singleVibratorComposedEffects_runsDifferentVibrations() {
         FakeVibratorControllerProvider fakeVibrator = mVibratorProviders.get(VIBRATOR_ID);
         fakeVibrator.setSupportedEffects(VibrationEffect.EFFECT_CLICK);
-        fakeVibrator.setSupportedPrimitives(
-                VibrationEffect.Composition.PRIMITIVE_CLICK,
-                VibrationEffect.Composition.PRIMITIVE_TICK);
+        fakeVibrator.setSupportedPrimitives(PRIMITIVE_CLICK, PRIMITIVE_TICK);
         fakeVibrator.setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS,
                 IVibrator.CAP_COMPOSE_PWLE_EFFECTS, IVibrator.CAP_AMPLITUDE_CONTROL);
         fakeVibrator.setMinFrequency(100);
@@ -820,8 +845,8 @@ public class VibrationThreadTest {
 
         VibrationEffect effect = VibrationEffect.startComposition()
                 .addEffect(VibrationEffect.createOneShot(10, 100))
-                .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 1f)
-                .addPrimitive(VibrationEffect.Composition.PRIMITIVE_TICK, 0.5f)
+                .addPrimitive(PRIMITIVE_CLICK, 1f)
+                .addPrimitive(PRIMITIVE_TICK, 0.5f)
                 .addEffect(VibrationEffect.get(VibrationEffect.EFFECT_CLICK))
                 .addEffect(VibrationEffect.startWaveform()
                         .addTransition(Duration.ofMillis(10),
@@ -836,13 +861,14 @@ public class VibrationThreadTest {
         // Use first duration the vibrator is turned on since we cannot estimate the clicks.
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(10L));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
-        verify(mControllerCallbacks, times(5)).onComplete(eq(VIBRATOR_ID), eq(vibration.id));
+        verify(mControllerCallbacks, times(5))
+                .onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.FINISHED);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
         assertEquals(Arrays.asList(
                 expectedOneShot(10),
-                expectedPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 1, 0),
-                expectedPrimitive(VibrationEffect.Composition.PRIMITIVE_TICK, 0.5f, 0),
+                expectedPrimitive(PRIMITIVE_CLICK, 1, 0),
+                expectedPrimitive(PRIMITIVE_TICK, 0.5f, 0),
                 expectedPrebaked(VibrationEffect.EFFECT_CLICK),
                 expectedRamp(/* startAmplitude= */ 0, /* endAmplitude= */ 0.5f,
                         /* startFrequencyHz= */ 150, /* endFrequencyHz= */ 100, /* duration= */ 10),
@@ -857,17 +883,15 @@ public class VibrationThreadTest {
     public void vibrate_singleVibratorComposedWithFallback_replacedInTheMiddleOfComposition() {
         FakeVibratorControllerProvider fakeVibrator = mVibratorProviders.get(VIBRATOR_ID);
         fakeVibrator.setSupportedEffects(VibrationEffect.EFFECT_CLICK);
-        fakeVibrator.setSupportedPrimitives(
-                VibrationEffect.Composition.PRIMITIVE_CLICK,
-                VibrationEffect.Composition.PRIMITIVE_TICK);
+        fakeVibrator.setSupportedPrimitives(PRIMITIVE_CLICK, PRIMITIVE_TICK);
         fakeVibrator.setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS);
 
         VibrationEffect fallback = VibrationEffect.createOneShot(10, 100);
         VibrationEffect effect = VibrationEffect.startComposition()
                 .addEffect(VibrationEffect.get(VibrationEffect.EFFECT_CLICK))
-                .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 1f)
+                .addPrimitive(PRIMITIVE_CLICK, 1f)
                 .addEffect(VibrationEffect.get(VibrationEffect.EFFECT_TICK))
-                .addPrimitive(VibrationEffect.Composition.PRIMITIVE_TICK, 0.5f)
+                .addPrimitive(PRIMITIVE_TICK, 0.5f)
                 .compose();
         HalVibration vibration = createVibration(CombinedVibration.createParallel(effect));
         vibration.fillFallbacks(unused -> fallback);
@@ -877,7 +901,8 @@ public class VibrationThreadTest {
         // Use first duration the vibrator is turned on since we cannot estimate the clicks.
         verify(mManagerHooks).noteVibratorOn(eq(UID), anyLong());
         verify(mManagerHooks).noteVibratorOff(eq(UID));
-        verify(mControllerCallbacks, times(4)).onComplete(eq(VIBRATOR_ID), eq(vibration.id));
+        verify(mControllerCallbacks, times(4))
+                .onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.FINISHED);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
 
@@ -894,7 +919,7 @@ public class VibrationThreadTest {
     }
 
     @Test
-    @EnableFlags(android.os.vibrator.Flags.FLAG_NORMALIZED_PWLE_EFFECTS)
+    @EnableFlags(Flags.FLAG_NORMALIZED_PWLE_EFFECTS)
     public void vibrate_singleVibratorPwle_runsComposePwleV2() {
         FakeVibratorControllerProvider fakeVibrator = mVibratorProviders.get(VIBRATOR_ID);
         fakeVibrator.setCapabilities(IVibrator.CAP_COMPOSE_PWLE_EFFECTS_V2);
@@ -915,7 +940,7 @@ public class VibrationThreadTest {
 
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(100L));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
-        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id));
+        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.FINISHED);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
         assertEquals(Arrays.asList(
@@ -929,7 +954,7 @@ public class VibrationThreadTest {
     }
 
     @Test
-    @EnableFlags(android.os.vibrator.Flags.FLAG_NORMALIZED_PWLE_EFFECTS)
+    @EnableFlags(Flags.FLAG_NORMALIZED_PWLE_EFFECTS)
     public void vibrate_singleVibratorBasicPwle_runsComposePwleV2() {
         FakeVibratorControllerProvider fakeVibrator = mVibratorProviders.get(VIBRATOR_ID);
         fakeVibrator.setCapabilities(IVibrator.CAP_COMPOSE_PWLE_EFFECTS_V2);
@@ -951,7 +976,7 @@ public class VibrationThreadTest {
 
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(220L));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
-        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id));
+        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.FINISHED);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
         assertEquals(Arrays.asList(
@@ -964,7 +989,7 @@ public class VibrationThreadTest {
     }
 
     @Test
-    @EnableFlags(android.os.vibrator.Flags.FLAG_NORMALIZED_PWLE_EFFECTS)
+    @EnableFlags(Flags.FLAG_NORMALIZED_PWLE_EFFECTS)
     public void vibrate_singleVibratorPwle_withInitialFrequency_runsComposePwleV2() {
         FakeVibratorControllerProvider fakeVibrator = mVibratorProviders.get(VIBRATOR_ID);
         fakeVibrator.setCapabilities(IVibrator.CAP_COMPOSE_PWLE_EFFECTS_V2);
@@ -987,7 +1012,7 @@ public class VibrationThreadTest {
 
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(100L));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
-        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id));
+        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.FINISHED);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
         assertEquals(Arrays.asList(
@@ -1001,7 +1026,7 @@ public class VibrationThreadTest {
     }
 
     @Test
-    @EnableFlags(android.os.vibrator.Flags.FLAG_NORMALIZED_PWLE_EFFECTS)
+    @EnableFlags(Flags.FLAG_NORMALIZED_PWLE_EFFECTS)
     public void vibrate_singleVibratorPwle_TooManyControlPoints_splitsAndRunsComposePwleV2() {
         FakeVibratorControllerProvider fakeVibrator = mVibratorProviders.get(VIBRATOR_ID);
         fakeVibrator.setCapabilities(IVibrator.CAP_COMPOSE_PWLE_EFFECTS_V2);
@@ -1027,7 +1052,8 @@ public class VibrationThreadTest {
         verifyCallbacksTriggered(vibration, Status.FINISHED);
         // Vibrator compose called 3 times with 2 segments instead of 2 times with 3 segments.
         // Using best split points instead of max-packing PWLEs.
-        verify(mControllerCallbacks, times(3)).onComplete(eq(VIBRATOR_ID), eq(vibration.id));
+        verify(mControllerCallbacks, times(3))
+                .onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
 
         assertEquals(Arrays.asList(
@@ -1043,7 +1069,7 @@ public class VibrationThreadTest {
     }
 
     @Test
-    @DisableFlags(android.os.vibrator.Flags.FLAG_NORMALIZED_PWLE_EFFECTS)
+    @DisableFlags(Flags.FLAG_NORMALIZED_PWLE_EFFECTS)
     public void vibrate_singleVibratorPwle_runsComposePwle() {
         FakeVibratorControllerProvider fakeVibrator = mVibratorProviders.get(VIBRATOR_ID);
         fakeVibrator.setCapabilities(IVibrator.CAP_COMPOSE_PWLE_EFFECTS);
@@ -1066,7 +1092,7 @@ public class VibrationThreadTest {
 
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(100L));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
-        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id));
+        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.FINISHED);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
         assertEquals(Arrays.asList(
@@ -1109,7 +1135,8 @@ public class VibrationThreadTest {
 
         // Vibrator compose called 3 times with 2 segments instead of 2 times with 3 segments.
         // Using best split points instead of max-packing PWLEs.
-        verify(mControllerCallbacks, times(3)).onComplete(eq(VIBRATOR_ID), eq(vibration.id));
+        verify(mControllerCallbacks, times(3))
+                .onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         assertEquals(6, fakeVibrator.getEffectSegments(vibration.id).size());
     }
 
@@ -1160,8 +1187,8 @@ public class VibrationThreadTest {
 
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(20L));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
-        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id));
-        verify(mControllerCallbacks, never()).onComplete(eq(2), eq(vibration.id));
+        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
+        verify(mControllerCallbacks, never()).onComplete(eq(2), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.FINISHED);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
 
@@ -1183,9 +1210,9 @@ public class VibrationThreadTest {
 
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(20L));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
-        verify(mControllerCallbacks).onComplete(eq(1), eq(vibration.id));
-        verify(mControllerCallbacks).onComplete(eq(2), eq(vibration.id));
-        verify(mControllerCallbacks).onComplete(eq(3), eq(vibration.id));
+        verify(mControllerCallbacks).onComplete(eq(1), eq(vibration.id), anyLong());
+        verify(mControllerCallbacks).onComplete(eq(2), eq(vibration.id), anyLong());
+        verify(mControllerCallbacks).onComplete(eq(3), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.FINISHED);
         assertFalse(mControllers.get(1).isVibrating());
         assertFalse(mControllers.get(2).isVibrating());
@@ -1207,11 +1234,10 @@ public class VibrationThreadTest {
         mVibratorProviders.get(2).setCapabilities(IVibrator.CAP_AMPLITUDE_CONTROL);
         mVibratorProviders.get(3).setCapabilities(IVibrator.CAP_AMPLITUDE_CONTROL);
         mVibratorProviders.get(4).setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS);
-        mVibratorProviders.get(4).setSupportedPrimitives(
-                VibrationEffect.Composition.PRIMITIVE_CLICK);
+        mVibratorProviders.get(4).setSupportedPrimitives(PRIMITIVE_CLICK);
 
         VibrationEffect composed = VibrationEffect.startComposition()
-                .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK)
+                .addPrimitive(PRIMITIVE_CLICK)
                 .compose();
         CombinedVibration effect = CombinedVibration.startParallel()
                 .addVibrator(1, VibrationEffect.get(VibrationEffect.EFFECT_CLICK))
@@ -1225,10 +1251,10 @@ public class VibrationThreadTest {
 
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(20L));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
-        verify(mControllerCallbacks).onComplete(eq(1), eq(vibration.id));
-        verify(mControllerCallbacks).onComplete(eq(2), eq(vibration.id));
-        verify(mControllerCallbacks).onComplete(eq(3), eq(vibration.id));
-        verify(mControllerCallbacks).onComplete(eq(4), eq(vibration.id));
+        verify(mControllerCallbacks).onComplete(eq(1), eq(vibration.id), anyLong());
+        verify(mControllerCallbacks).onComplete(eq(2), eq(vibration.id), anyLong());
+        verify(mControllerCallbacks).onComplete(eq(3), eq(vibration.id), anyLong());
+        verify(mControllerCallbacks).onComplete(eq(4), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.FINISHED);
         assertFalse(mControllers.get(1).isVibrating());
         assertFalse(mControllers.get(2).isVibrating());
@@ -1243,8 +1269,7 @@ public class VibrationThreadTest {
         assertEquals(Arrays.asList(expectedOneShot(20)),
                 mVibratorProviders.get(3).getEffectSegments(vibration.id));
         assertEquals(expectedAmplitudes(1, 2), mVibratorProviders.get(3).getAmplitudes());
-        assertEquals(Arrays.asList(
-                expectedPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 1, 0)),
+        assertEquals(Arrays.asList(expectedPrimitive(PRIMITIVE_CLICK, 1, 0)),
                 mVibratorProviders.get(4).getEffectSegments(vibration.id));
     }
 
@@ -1253,12 +1278,11 @@ public class VibrationThreadTest {
         mockVibrators(1, 2, 3);
         mVibratorProviders.get(1).setCapabilities(IVibrator.CAP_AMPLITUDE_CONTROL);
         mVibratorProviders.get(2).setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS);
-        mVibratorProviders.get(2).setSupportedPrimitives(
-                VibrationEffect.Composition.PRIMITIVE_CLICK);
+        mVibratorProviders.get(2).setSupportedPrimitives(PRIMITIVE_CLICK);
         mVibratorProviders.get(3).setSupportedEffects(VibrationEffect.EFFECT_CLICK);
 
         VibrationEffect composed = VibrationEffect.startComposition()
-                .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK)
+                .addPrimitive(PRIMITIVE_CLICK)
                 .compose();
         CombinedVibration effect = CombinedVibration.startSequential()
                 .addNext(3, VibrationEffect.get(VibrationEffect.EFFECT_CLICK), /* delay= */ 50)
@@ -1268,10 +1292,10 @@ public class VibrationThreadTest {
         HalVibration vibration = startThreadAndDispatcher(effect);
 
         waitForCompletion();
-        InOrder controllerVerifier = inOrder(mControllerCallbacks);
-        controllerVerifier.verify(mControllerCallbacks).onComplete(eq(3), eq(vibration.id));
-        controllerVerifier.verify(mControllerCallbacks).onComplete(eq(1), eq(vibration.id));
-        controllerVerifier.verify(mControllerCallbacks).onComplete(eq(2), eq(vibration.id));
+        InOrder verifier = inOrder(mControllerCallbacks);
+        verifier.verify(mControllerCallbacks).onComplete(eq(3), eq(vibration.id), anyLong());
+        verifier.verify(mControllerCallbacks).onComplete(eq(1), eq(vibration.id), anyLong());
+        verifier.verify(mControllerCallbacks).onComplete(eq(2), eq(vibration.id), anyLong());
 
         InOrder batteryVerifier = inOrder(mManagerHooks);
         batteryVerifier.verify(mManagerHooks).noteVibratorOn(eq(UID), eq(20L));
@@ -1289,8 +1313,7 @@ public class VibrationThreadTest {
         assertEquals(Arrays.asList(expectedOneShot(10)),
                 mVibratorProviders.get(1).getEffectSegments(vibration.id));
         assertEquals(expectedAmplitudes(100), mVibratorProviders.get(1).getAmplitudes());
-        assertEquals(Arrays.asList(
-                expectedPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 1, 0)),
+        assertEquals(Arrays.asList(expectedPrimitive(PRIMITIVE_CLICK, 1, 0)),
                 mVibratorProviders.get(2).getEffectSegments(vibration.id));
         assertEquals(Arrays.asList(expectedPrebaked(VibrationEffect.EFFECT_CLICK)),
                 mVibratorProviders.get(3).getEffectSegments(vibration.id));
@@ -1301,15 +1324,13 @@ public class VibrationThreadTest {
         int[] vibratorIds = new int[]{1, 2};
         mockVibrators(vibratorIds);
         mVibratorProviders.get(1).setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS);
-        mVibratorProviders.get(1).setSupportedPrimitives(
-                VibrationEffect.Composition.PRIMITIVE_CLICK);
+        mVibratorProviders.get(1).setSupportedPrimitives(PRIMITIVE_CLICK);
         mVibratorProviders.get(2).setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS);
-        mVibratorProviders.get(2).setSupportedPrimitives(
-                VibrationEffect.Composition.PRIMITIVE_CLICK);
+        mVibratorProviders.get(2).setSupportedPrimitives(PRIMITIVE_CLICK);
         when(mManagerHooks.prepareSyncedVibration(anyLong(), eq(vibratorIds))).thenReturn(true);
 
         VibrationEffect composed = VibrationEffect.startComposition()
-                .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 1, 100)
+                .addPrimitive(PRIMITIVE_CLICK, 1, 100)
                 .compose();
         CombinedVibration effect = CombinedVibration.createParallel(composed);
         // We create the HalVibration here to obtain the vibration id and use it to mock the
@@ -1331,8 +1352,7 @@ public class VibrationThreadTest {
         verify(mManagerHooks, never()).cancelSyncedVibration();
         verifyCallbacksTriggered(vibration, Status.FINISHED);
 
-        VibrationEffectSegment expected = expectedPrimitive(
-                VibrationEffect.Composition.PRIMITIVE_CLICK, 1, 100);
+        VibrationEffectSegment expected = expectedPrimitive(PRIMITIVE_CLICK, 1, 100);
         assertEquals(Arrays.asList(expected),
                 mVibratorProviders.get(1).getEffectSegments(vibration.id));
         assertEquals(Arrays.asList(expected),
@@ -1345,12 +1365,11 @@ public class VibrationThreadTest {
         mockVibrators(vibratorIds);
         mVibratorProviders.get(1).setSupportedEffects(VibrationEffect.EFFECT_CLICK);
         mVibratorProviders.get(4).setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS);
-        mVibratorProviders.get(4).setSupportedPrimitives(
-                VibrationEffect.Composition.PRIMITIVE_CLICK);
+        mVibratorProviders.get(4).setSupportedPrimitives(PRIMITIVE_CLICK);
         when(mManagerHooks.prepareSyncedVibration(anyLong(), any())).thenReturn(true);
 
         VibrationEffect composed = VibrationEffect.startComposition()
-                .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK)
+                .addPrimitive(PRIMITIVE_CLICK)
                 .compose();
         CombinedVibration effect = CombinedVibration.startParallel()
                 .addVibrator(1, VibrationEffect.get(VibrationEffect.EFFECT_CLICK))
@@ -1463,9 +1482,9 @@ public class VibrationThreadTest {
 
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(80L));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
-        verify(mControllerCallbacks).onComplete(eq(1), eq(vibration.id));
-        verify(mControllerCallbacks).onComplete(eq(2), eq(vibration.id));
-        verify(mControllerCallbacks).onComplete(eq(3), eq(vibration.id));
+        verify(mControllerCallbacks).onComplete(eq(1), eq(vibration.id), anyLong());
+        verify(mControllerCallbacks).onComplete(eq(2), eq(vibration.id), anyLong());
+        verify(mControllerCallbacks).onComplete(eq(3), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.FINISHED);
         assertFalse(mControllers.get(1).isVibrating());
         assertFalse(mControllers.get(2).isVibrating());
@@ -1505,7 +1524,7 @@ public class VibrationThreadTest {
         waitForCompletion(rampDownDuration + TEST_TIMEOUT_MILLIS);
         long completionTime = SystemClock.elapsedRealtime();
 
-        verify(mControllerCallbacks).onComplete(VIBRATOR_ID, vibration.id);
+        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         // Vibration ends after duration, thread completed after ramp down
         assertThat(vibrationEndTime - startTime).isAtLeast(expectedDuration);
         assertThat(vibrationEndTime - startTime).isLessThan(expectedDuration + rampDownDuration);
@@ -1534,7 +1553,7 @@ public class VibrationThreadTest {
 
         waitForCompletion(TEST_TIMEOUT_MILLIS);
 
-        verify(mControllerCallbacks).onComplete(VIBRATOR_ID, vibration.id);
+        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         assertThat(vibrationEndTime - startTime).isAtLeast(expectedDuration + callbackDelay);
     }
 
@@ -1563,7 +1582,8 @@ public class VibrationThreadTest {
         waitForCompletion(callbackDelay + TEST_TIMEOUT_MILLIS);
         long completionTime = SystemClock.elapsedRealtime();
 
-        verify(mControllerCallbacks, never()).onComplete(VIBRATOR_ID, vibration.id);
+        verify(mControllerCallbacks, never())
+                .onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         // Vibration ends and thread completes after timeout, before the HAL callback
         assertThat(vibrationEndTime - startTime).isAtLeast(expectedDuration + callbackTimeout);
         assertThat(vibrationEndTime - startTime).isLessThan(expectedDuration + callbackDelay);
@@ -1639,15 +1659,14 @@ public class VibrationThreadTest {
         mockVibrators(1, 2);
         mVibratorProviders.get(1).setSupportedEffects(VibrationEffect.EFFECT_CLICK);
         mVibratorProviders.get(2).setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS);
-        mVibratorProviders.get(2).setSupportedPrimitives(
-                VibrationEffect.Composition.PRIMITIVE_CLICK);
+        mVibratorProviders.get(2).setSupportedPrimitives(PRIMITIVE_CLICK);
 
         CombinedVibration effect = CombinedVibration.startParallel()
                 .addVibrator(1, VibrationEffect.get(VibrationEffect.EFFECT_CLICK))
                 .addVibrator(2, VibrationEffect.startComposition()
-                        .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 1f, 100)
-                        .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 1f, 100)
-                        .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 1f, 100)
+                        .addPrimitive(PRIMITIVE_CLICK, 1f, 100)
+                        .addPrimitive(PRIMITIVE_CLICK, 1f, 100)
+                        .addPrimitive(PRIMITIVE_CLICK, 1f, 100)
                         .compose())
                 .combine();
         HalVibration vibration = startThreadAndDispatcher(effect);
@@ -1672,7 +1691,7 @@ public class VibrationThreadTest {
     }
 
     @Test
-    @RequiresFlagsEnabled(android.os.vibrator.Flags.FLAG_VENDOR_VIBRATION_EFFECTS)
+    @EnableFlags(Flags.FLAG_VENDOR_VIBRATION_EFFECTS)
     public void vibrate_multipleVendorEffectCancel_cancelsVibrationImmediately() throws Exception {
         mockVibrators(1, 2);
         mVibratorProviders.get(1).setCapabilities(IVibrator.CAP_PERFORM_VENDOR_EFFECTS);
@@ -1767,7 +1786,7 @@ public class VibrationThreadTest {
         HalVibration vibration = startThreadAndDispatcher(effect);
         waitForCompletion();
 
-        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id));
+        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.FINISHED);
 
         // Duration extended for 5 + 5 + 5 + 15.
@@ -1847,7 +1866,7 @@ public class VibrationThreadTest {
         HalVibration vibration = startThreadAndDispatcher(effect);
         waitForCompletion();
 
-        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id));
+        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.FINISHED);
 
         assertEquals(Arrays.asList(expectedPrebaked(VibrationEffect.EFFECT_CLICK)),
@@ -1856,7 +1875,7 @@ public class VibrationThreadTest {
     }
 
     @Test
-    @RequiresFlagsEnabled(android.os.vibrator.Flags.FLAG_VENDOR_VIBRATION_EFFECTS)
+    @EnableFlags(Flags.FLAG_VENDOR_VIBRATION_EFFECTS)
     public void vibrate_vendorEffectWithRampDown_doesNotAddRampDown() {
         when(mVibrationConfigMock.getRampDownDurationMs()).thenReturn(15);
         mVibratorProviders.get(VIBRATOR_ID).setCapabilities(IVibrator.CAP_PERFORM_VENDOR_EFFECTS);
@@ -1865,7 +1884,7 @@ public class VibrationThreadTest {
         HalVibration vibration = startThreadAndDispatcher(effect);
         waitForCompletion();
 
-        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id));
+        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.FINISHED);
 
         assertThat(mVibratorProviders.get(VIBRATOR_ID).getVendorEffects(vibration.id))
@@ -1879,26 +1898,24 @@ public class VibrationThreadTest {
         when(mVibrationConfigMock.getRampDownDurationMs()).thenReturn(15);
         mVibratorProviders.get(VIBRATOR_ID).setCapabilities(IVibrator.CAP_AMPLITUDE_CONTROL,
                 IVibrator.CAP_COMPOSE_EFFECTS);
-        mVibratorProviders.get(VIBRATOR_ID).setSupportedPrimitives(
-                VibrationEffect.Composition.PRIMITIVE_CLICK);
+        mVibratorProviders.get(VIBRATOR_ID).setSupportedPrimitives(PRIMITIVE_CLICK);
 
         VibrationEffect effect = VibrationEffect.startComposition()
-                .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK)
+                .addPrimitive(PRIMITIVE_CLICK)
                 .compose();
         HalVibration vibration = startThreadAndDispatcher(effect);
         waitForCompletion();
 
-        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id));
+        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.FINISHED);
 
-        assertEquals(
-                Arrays.asList(expectedPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 1, 0)),
+        assertEquals(Arrays.asList(expectedPrimitive(PRIMITIVE_CLICK, 1, 0)),
                 mVibratorProviders.get(VIBRATOR_ID).getEffectSegments(vibration.id));
         assertTrue(mVibratorProviders.get(VIBRATOR_ID).getAmplitudes().isEmpty());
     }
 
     @Test
-    @DisableFlags(android.os.vibrator.Flags.FLAG_NORMALIZED_PWLE_EFFECTS)
+    @DisableFlags(Flags.FLAG_NORMALIZED_PWLE_EFFECTS)
     public void vibrate_pwleWithRampDown_doesNotAddRampDown() {
         when(mVibrationConfigMock.getRampDownDurationMs()).thenReturn(15);
         FakeVibratorControllerProvider fakeVibrator = mVibratorProviders.get(VIBRATOR_ID);
@@ -1916,7 +1933,7 @@ public class VibrationThreadTest {
         HalVibration vibration = startThreadAndDispatcher(effect);
         waitForCompletion();
 
-        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id));
+        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration.id), anyLong());
         verifyCallbacksTriggered(vibration, Status.FINISHED);
 
         assertEquals(Arrays.asList(expectedRamp(0, 1, 150, 150, 1)),
@@ -1928,8 +1945,7 @@ public class VibrationThreadTest {
     public void vibrate_multipleVibrations_withCancel() throws Exception {
         mVibratorProviders.get(VIBRATOR_ID).setSupportedEffects(
                 VibrationEffect.EFFECT_CLICK, VibrationEffect.EFFECT_TICK);
-        mVibratorProviders.get(VIBRATOR_ID).setSupportedPrimitives(
-                VibrationEffect.Composition.PRIMITIVE_CLICK);
+        mVibratorProviders.get(VIBRATOR_ID).setSupportedPrimitives(PRIMITIVE_CLICK);
         mVibratorProviders.get(VIBRATOR_ID).setCapabilities(IVibrator.CAP_AMPLITUDE_CONTROL,
                 IVibrator.CAP_COMPOSE_EFFECTS);
 
@@ -1940,7 +1956,7 @@ public class VibrationThreadTest {
                 .repeatEffectIndefinitely(VibrationEffect.get(VibrationEffect.EFFECT_TICK))
                 .compose();
         VibrationEffect effect3 = VibrationEffect.startComposition()
-                .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK)
+                .addPrimitive(PRIMITIVE_CLICK)
                 .compose();
         VibrationEffect effect4 = VibrationEffect.createOneShot(8000, 100);
         VibrationEffect effect5 = VibrationEffect.get(VibrationEffect.EFFECT_CLICK);
@@ -1974,14 +1990,15 @@ public class VibrationThreadTest {
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
 
         // Effect1
-        verify(mControllerCallbacks).onComplete(VIBRATOR_ID, vibration1.id);
+        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration1.id), anyLong());
         verifyCallbacksTriggered(vibration1, Status.FINISHED);
 
         assertEquals(Arrays.asList(expectedPrebaked(VibrationEffect.EFFECT_CLICK)),
                 fakeVibrator.getEffectSegments(vibration1.id));
 
         // Effect2: repeating, cancelled.
-        verify(mControllerCallbacks, atLeast(2)).onComplete(VIBRATOR_ID, vibration2.id);
+        verify(mControllerCallbacks, atLeast(2))
+                .onComplete(eq(VIBRATOR_ID), eq(vibration2.id), anyLong());
         verifyCallbacksTriggered(vibration2, Status.CANCELLED_BY_USER);
 
         // The exact count of segments might vary, so just check that there's more than 2 and
@@ -1994,10 +2011,9 @@ public class VibrationThreadTest {
         }
 
         // Effect3
-        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration3.id));
+        verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibration3.id), anyLong());
         verifyCallbacksTriggered(vibration3, Status.FINISHED);
-        assertEquals(Arrays.asList(
-                        expectedPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 1, 0)),
+        assertEquals(Arrays.asList(expectedPrimitive(PRIMITIVE_CLICK, 1, 0)),
                 fakeVibrator.getEffectSegments(vibration3.id));
 
         // Effect4: cancelled quickly.
@@ -2014,8 +2030,7 @@ public class VibrationThreadTest {
         mockVibrators(1, 2, 3);
         mVibratorProviders.get(1).setCapabilities(IVibrator.CAP_AMPLITUDE_CONTROL);
         mVibratorProviders.get(2).setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS);
-        mVibratorProviders.get(2).setSupportedPrimitives(
-                VibrationEffect.Composition.PRIMITIVE_CLICK);
+        mVibratorProviders.get(2).setSupportedPrimitives(PRIMITIVE_CLICK);
         mVibratorProviders.get(3).setSupportedEffects(VibrationEffect.EFFECT_CLICK);
 
         CombinedVibration effect = CombinedVibration.startSequential()
@@ -2029,8 +2044,7 @@ public class VibrationThreadTest {
                         /* delay= */ TEST_TIMEOUT_MILLIS)
                 .addNext(2,
                         VibrationEffect.startComposition()
-                                .addPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 1,
-                                        /* delay= */ TEST_TIMEOUT_MILLIS)
+                                .addPrimitive(PRIMITIVE_CLICK, 1, /* delay= */ TEST_TIMEOUT_MILLIS)
                                 .compose(),
                         /* delay= */ TEST_TIMEOUT_MILLIS)
                 .combine();
@@ -2051,8 +2065,7 @@ public class VibrationThreadTest {
         assertEquals(Arrays.asList(expectedOneShot(TEST_TIMEOUT_MILLIS)),
                 mVibratorProviders.get(1).getEffectSegments(vibration.id));
         assertEquals(expectedAmplitudes(255), mVibratorProviders.get(1).getAmplitudes());
-        assertEquals(Arrays.asList(expectedPrimitive(
-                VibrationEffect.Composition.PRIMITIVE_CLICK, 1, TEST_TIMEOUT_MILLIS)),
+        assertEquals(Arrays.asList(expectedPrimitive(PRIMITIVE_CLICK, 1, TEST_TIMEOUT_MILLIS)),
                 mVibratorProviders.get(2).getEffectSegments(vibration.id));
         assertEquals(Arrays.asList(expectedPrebaked(VibrationEffect.EFFECT_CLICK)),
                 mVibratorProviders.get(3).getEffectSegments(vibration.id));
@@ -2205,7 +2218,7 @@ public class VibrationThreadTest {
 
     private void verifyCallbacksTriggered(HalVibration vibration, Status expectedStatus) {
         assertThat(vibration.getStatus()).isEqualTo(expectedStatus);
-        verify(mManagerHooks).onVibrationThreadReleased(vibration.id);
+        verify(mManagerHooks).onVibrationThreadReleased(eq(vibration.id));
     }
 
     private static final class TestLooperAutoDispatcher extends Thread {
diff --git a/services/tests/vibrator/src/com/android/server/vibrator/VibratorControllerTest.java b/services/tests/vibrator/src/com/android/server/vibrator/VibratorControllerTest.java
index 0978f48491cc..4df13deaed7d 100644
--- a/services/tests/vibrator/src/com/android/server/vibrator/VibratorControllerTest.java
+++ b/services/tests/vibrator/src/com/android/server/vibrator/VibratorControllerTest.java
@@ -203,9 +203,10 @@ public class VibratorControllerTest {
 
     @Test
     public void setAmplitude_vibratorVibrating_setsAmplitude() {
-        when(mNativeWrapperMock.on(anyLong(), anyLong())).thenAnswer(args -> args.getArgument(0));
+        when(mNativeWrapperMock.on(anyLong(), anyLong(), anyLong()))
+                .thenAnswer(args -> args.getArgument(0));
         VibratorController controller = createController();
-        controller.on(100, /* vibrationId= */ 1);
+        controller.on(100, 1, 1);
         assertTrue(controller.isVibrating());
         assertEquals(-1, controller.getCurrentAmplitude(), /* delta= */ 0);
 
@@ -215,81 +216,84 @@ public class VibratorControllerTest {
 
     @Test
     public void on_withDuration_turnsVibratorOn() {
-        when(mNativeWrapperMock.on(anyLong(), anyLong())).thenAnswer(args -> args.getArgument(0));
+        when(mNativeWrapperMock.on(anyLong(), anyLong(), anyLong()))
+                .thenAnswer(args -> args.getArgument(0));
         VibratorController controller = createController();
-        controller.on(100, 10);
+        controller.on(100, 10, 20);
 
         assertTrue(controller.isVibrating());
-        verify(mNativeWrapperMock).on(eq(100L), eq(10L));
+        verify(mNativeWrapperMock).on(eq(100L), eq(10L), eq(20L));
     }
 
     @Test
     public void on_withPrebaked_performsEffect() {
-        when(mNativeWrapperMock.perform(anyLong(), anyLong(), anyLong())).thenReturn(10L);
+        when(mNativeWrapperMock.perform(anyLong(), anyLong(), anyLong(), anyLong()))
+                .thenReturn(10L);
         VibratorController controller = createController();
 
         PrebakedSegment prebaked = createPrebaked(VibrationEffect.EFFECT_CLICK,
                 VibrationEffect.EFFECT_STRENGTH_MEDIUM);
-        assertEquals(10L, controller.on(prebaked, 11));
+        assertEquals(10L, controller.on(prebaked, 11, 23));
 
         assertTrue(controller.isVibrating());
         verify(mNativeWrapperMock).perform(eq((long) VibrationEffect.EFFECT_CLICK),
-                eq((long) VibrationEffect.EFFECT_STRENGTH_MEDIUM), eq(11L));
+                eq((long) VibrationEffect.EFFECT_STRENGTH_MEDIUM), eq(11L), eq(23L));
     }
 
     @Test
     public void on_withComposed_performsEffect() {
         mockVibratorCapabilities(IVibrator.CAP_COMPOSE_EFFECTS);
-        when(mNativeWrapperMock.compose(any(), anyLong())).thenReturn(15L);
+        when(mNativeWrapperMock.compose(any(), anyLong(), anyLong())).thenReturn(15L);
         VibratorController controller = createController();
 
         PrimitiveSegment[] primitives = new PrimitiveSegment[]{
                 new PrimitiveSegment(VibrationEffect.Composition.PRIMITIVE_CLICK, 0.5f, 10)
         };
-        assertEquals(15L, controller.on(primitives, 12));
+        assertEquals(15L, controller.on(primitives, 12, 34));
 
         assertTrue(controller.isVibrating());
-        verify(mNativeWrapperMock).compose(eq(primitives), eq(12L));
+        verify(mNativeWrapperMock).compose(eq(primitives), eq(12L), eq(34L));
     }
 
     @Test
     public void on_withComposedPwle_performsEffect() {
         mockVibratorCapabilities(IVibrator.CAP_COMPOSE_PWLE_EFFECTS);
-        when(mNativeWrapperMock.composePwle(any(), anyInt(), anyLong())).thenReturn(15L);
+        when(mNativeWrapperMock.composePwle(any(), anyInt(), anyLong(), anyLong())).thenReturn(15L);
         VibratorController controller = createController();
 
         RampSegment[] primitives = new RampSegment[]{
                 new RampSegment(/* startAmplitude= */ 0, /* endAmplitude= */ 1,
                         /* startFrequencyHz= */ 100, /* endFrequencyHz= */ 200, /* duration= */ 10)
         };
-        assertEquals(15L, controller.on(primitives, 12));
+        assertEquals(15L, controller.on(primitives, 12, 45));
         assertTrue(controller.isVibrating());
 
-        verify(mNativeWrapperMock).composePwle(eq(primitives), eq(Braking.NONE), eq(12L));
+        verify(mNativeWrapperMock).composePwle(eq(primitives), eq(Braking.NONE), eq(12L), eq(45L));
     }
 
     @Test
     public void on_withComposedPwleV2_performsEffect() {
         mockVibratorCapabilities(IVibrator.CAP_COMPOSE_PWLE_EFFECTS_V2);
-        when(mNativeWrapperMock.composePwleV2(any(), anyLong())).thenReturn(15L);
+        when(mNativeWrapperMock.composePwleV2(any(), anyLong(), anyLong())).thenReturn(15L);
         VibratorController controller = createController();
 
         PwlePoint[] primitives = new PwlePoint[]{
                 new PwlePoint(/*amplitude=*/ 0, /*frequencyHz=*/ 100, /*timeMillis=*/ 0),
                 new PwlePoint(/*amplitude=*/ 1, /*frequencyHz=*/ 200, /*timeMillis=*/ 10)
         };
-        assertEquals(15L, controller.on(primitives, 12));
+        assertEquals(15L, controller.on(primitives, 12, 53));
         assertTrue(controller.isVibrating());
 
-        verify(mNativeWrapperMock).composePwleV2(eq(primitives), eq(12L));
+        verify(mNativeWrapperMock).composePwleV2(eq(primitives), eq(12L), eq(53L));
     }
 
     @Test
     public void off_turnsOffVibrator() {
-        when(mNativeWrapperMock.on(anyLong(), anyLong())).thenAnswer(args -> args.getArgument(0));
+        when(mNativeWrapperMock.on(anyLong(), anyLong(), anyLong()))
+                .thenAnswer(args -> args.getArgument(0));
         VibratorController controller = createController();
 
-        controller.on(100, 1);
+        controller.on(100, 1, 1);
         assertTrue(controller.isVibrating());
 
         controller.off();
@@ -301,10 +305,11 @@ public class VibratorControllerTest {
     @Test
     public void reset_turnsOffVibratorAndDisablesExternalControl() {
         mockVibratorCapabilities(IVibrator.CAP_EXTERNAL_CONTROL);
-        when(mNativeWrapperMock.on(anyLong(), anyLong())).thenAnswer(args -> args.getArgument(0));
+        when(mNativeWrapperMock.on(anyLong(), anyLong(), anyLong()))
+                .thenAnswer(args -> args.getArgument(0));
         VibratorController controller = createController();
 
-        controller.on(100, 1);
+        controller.on(100, 1, 1);
         assertTrue(controller.isVibrating());
 
         controller.reset();
@@ -315,12 +320,13 @@ public class VibratorControllerTest {
 
     @Test
     public void registerVibratorStateListener_callbacksAreTriggered() throws Exception {
-        when(mNativeWrapperMock.on(anyLong(), anyLong())).thenAnswer(args -> args.getArgument(0));
+        when(mNativeWrapperMock.on(anyLong(), anyLong(), anyLong()))
+                .thenAnswer(args -> args.getArgument(0));
         VibratorController controller = createController();
 
         controller.registerVibratorStateListener(mVibratorStateListenerMock);
-        controller.on(10, 1);
-        controller.on(100, 2);
+        controller.on(10, 1, 1);
+        controller.on(100, 2, 1);
         controller.off();
         controller.off();
 
@@ -334,19 +340,20 @@ public class VibratorControllerTest {
 
     @Test
     public void unregisterVibratorStateListener_callbackNotTriggeredAfter() throws Exception {
-        when(mNativeWrapperMock.on(anyLong(), anyLong())).thenAnswer(args -> args.getArgument(0));
+        when(mNativeWrapperMock.on(anyLong(), anyLong(), anyLong()))
+                .thenAnswer(args -> args.getArgument(0));
         VibratorController controller = createController();
 
         controller.registerVibratorStateListener(mVibratorStateListenerMock);
         verify(mVibratorStateListenerMock).onVibrating(false);
 
-        controller.on(10, 1);
+        controller.on(10, 1, 1);
         verify(mVibratorStateListenerMock).onVibrating(true);
 
         controller.unregisterVibratorStateListener(mVibratorStateListenerMock);
         Mockito.clearInvocations(mVibratorStateListenerMock);
 
-        controller.on(10, 1);
+        controller.on(10, 1, 1);
         verifyNoMoreInteractions(mVibratorStateListenerMock);
     }
 
diff --git a/services/tests/vibrator/utils/com/android/server/vibrator/FakeVibratorControllerProvider.java b/services/tests/vibrator/utils/com/android/server/vibrator/FakeVibratorControllerProvider.java
index 3f3476716831..bd806b7f7f37 100644
--- a/services/tests/vibrator/utils/com/android/server/vibrator/FakeVibratorControllerProvider.java
+++ b/services/tests/vibrator/utils/com/android/server/vibrator/FakeVibratorControllerProvider.java
@@ -118,11 +118,11 @@ public final class FakeVibratorControllerProvider {
         }
 
         @Override
-        public long on(long milliseconds, long vibrationId) {
+        public long on(long milliseconds, long vibrationId, long stepId) {
             recordEffectSegment(vibrationId, new StepSegment(VibrationEffect.DEFAULT_AMPLITUDE,
                     /* frequencyHz= */ 0, (int) milliseconds));
             applyLatency(mOnLatency);
-            scheduleListener(milliseconds, vibrationId);
+            scheduleListener(milliseconds, vibrationId, stepId);
             return milliseconds;
         }
 
@@ -139,7 +139,7 @@ public final class FakeVibratorControllerProvider {
         }
 
         @Override
-        public long perform(long effect, long strength, long vibrationId) {
+        public long perform(long effect, long strength, long vibrationId, long stepId) {
             if (mSupportedEffects == null
                     || Arrays.binarySearch(mSupportedEffects, (int) effect) < 0) {
                 return 0;
@@ -147,13 +147,13 @@ public final class FakeVibratorControllerProvider {
             recordEffectSegment(vibrationId,
                     new PrebakedSegment((int) effect, false, (int) strength));
             applyLatency(mOnLatency);
-            scheduleListener(EFFECT_DURATION, vibrationId);
+            scheduleListener(EFFECT_DURATION, vibrationId, stepId);
             return EFFECT_DURATION;
         }
 
         @Override
         public long performVendorEffect(Parcel vendorData, long strength, float scale,
-                float adaptiveScale, long vibrationId) {
+                float adaptiveScale, long vibrationId, long stepId) {
             if ((mCapabilities & IVibrator.CAP_PERFORM_VENDOR_EFFECTS) == 0) {
                 return 0;
             }
@@ -161,13 +161,13 @@ public final class FakeVibratorControllerProvider {
             recordVendorEffect(vibrationId,
                     new VibrationEffect.VendorEffect(bundle, (int) strength, scale, adaptiveScale));
             applyLatency(mOnLatency);
-            scheduleListener(mVendorEffectDuration, vibrationId);
+            scheduleListener(mVendorEffectDuration, vibrationId, stepId);
             // HAL has unknown duration for vendor effects.
             return Long.MAX_VALUE;
         }
 
         @Override
-        public long compose(PrimitiveSegment[] primitives, long vibrationId) {
+        public long compose(PrimitiveSegment[] primitives, long vibrationId, long stepId) {
             if (mSupportedPrimitives == null) {
                 return 0;
             }
@@ -182,12 +182,13 @@ public final class FakeVibratorControllerProvider {
                 recordEffectSegment(vibrationId, primitive);
             }
             applyLatency(mOnLatency);
-            scheduleListener(duration, vibrationId);
+            scheduleListener(duration, vibrationId, stepId);
             return duration;
         }
 
         @Override
-        public long composePwle(RampSegment[] primitives, int braking, long vibrationId) {
+        public long composePwle(RampSegment[] primitives, int braking, long vibrationId,
+                long stepId) {
             long duration = 0;
             for (RampSegment primitive : primitives) {
                 duration += primitive.getDuration();
@@ -195,19 +196,19 @@ public final class FakeVibratorControllerProvider {
             }
             recordBraking(vibrationId, braking);
             applyLatency(mOnLatency);
-            scheduleListener(duration, vibrationId);
+            scheduleListener(duration, vibrationId, stepId);
             return duration;
         }
 
         @Override
-        public long composePwleV2(PwlePoint[] pwlePoints, long vibrationId) {
+        public long composePwleV2(PwlePoint[] pwlePoints, long vibrationId, long stepId) {
             long duration = 0;
             for (PwlePoint pwlePoint: pwlePoints) {
                 duration += pwlePoint.getTimeMillis();
                 recordEffectPwlePoint(vibrationId, pwlePoint);
             }
             applyLatency(mOnLatency);
-            scheduleListener(duration, vibrationId);
+            scheduleListener(duration, vibrationId, stepId);
 
             return duration;
         }
@@ -263,8 +264,8 @@ public final class FakeVibratorControllerProvider {
             }
         }
 
-        private void scheduleListener(long vibrationDuration, long vibrationId) {
-            mHandler.postDelayed(() -> listener.onComplete(vibratorId, vibrationId),
+        private void scheduleListener(long vibrationDuration, long vibrationId, long stepId) {
+            mHandler.postDelayed(() -> listener.onComplete(vibratorId, vibrationId, stepId),
                     vibrationDuration + mCompletionCallbackDelay);
         }
     }