simplify logic related to whether a manual acknowledgement was taken. Return null instead of throwing exception

Author: sbSteveK

src/main/java/software/amazon/awssdk/crt/mqtt5/PublishReturn.java

@@ -30,14 +30,6 @@ public class PublishReturn {
      */
     private long threadID;
 
-    /**
-     * Set to true when {@code acquirePublishAcknowledgementControl()} is called, indicating that
-     * the user has taken manual control of the publish acknowledgement. Native code reads this
-     * via {@code wasControlAcquired()} after the callback returns to decide whether to
-     * auto-invoke the publish acknowledgement.
-     */
-    private boolean controlAcquired;
-
     /**
      * Returns the PublishPacket returned from the server or Null if none was returned.
      * @return The PublishPacket returned from the server.
@@ -53,58 +45,29 @@ public PublishPacket getPublishPacket() {
      * at a later time to send the publish acknowledgement to the broker.
      *
      * <p><b>Important:</b> This method must be called within the
-     * {@link Mqtt5ClientOptions.PublishEvents#onMessageReceived} callback. Calling it after the
-     * callback returns will throw an {@link IllegalStateException}.</p>
+     * {@link Mqtt5ClientOptions.PublishEvents#onMessageReceived} callback. Calling it outside the
+     * callback (wrong thread) or after it has already been called will return {@code null}.</p>
      *
-     * <p>This method may only be called once per received PUBLISH. Subsequent calls will throw
-     * an {@link IllegalStateException}.</p>
+     * <p>This method may only be called once per received PUBLISH. Subsequent calls will return
+     * {@code null}.</p>
      *
      * <p>If this method is not called, the client will automatically send a publish acknowledgment
      * for QoS 1 messages when the callback returns.</p>
      *
-     * @return A {@link Mqtt5PublishAcknowledgementControlHandle} that can be used to manually send the acknowledgement.
-     * @throws IllegalStateException if called outside the onMessageReceived callback, called more than once,
-     *                               or called on a QoS 0 message.
+     * @return A {@link Mqtt5PublishAcknowledgementControlHandle} that can be used to manually send
+     *         the acknowledgement, or {@code null} if called outside the callback, called more than
+     *         once, or called on a QoS 0 message.
      */
     public synchronized Mqtt5PublishAcknowledgementControlHandle acquirePublishAcknowledgementControl() {
         if (controlId == 0 || threadID != Thread.currentThread().getId()) {
-            throw new IllegalStateException(
-                "acquirePublishAcknowledgementControl() must be called within the onMessageReceived callback and may only be called once.");
+            return null;
         }
         long acquiredControlId = controlId;
         /* Zero out so it can't be double-called */
         controlId = 0;
-        controlAcquired = true;
         return new Mqtt5PublishAcknowledgementControlHandle(acquiredControlId);
     }
 
-    /**
-     * Returns whether the user called {@link #acquirePublishAcknowledgementControl()} during the
-     * {@link Mqtt5ClientOptions.PublishEvents#onMessageReceived} callback.
-     *
-     * <p>This is called by native/JNI code after the callback returns to determine whether to
-     * automatically invoke the publish acknowledgement (PUBACK). If this returns {@code false},
-     * native code will auto-invoke the PUBACK; if {@code true}, the user is responsible for
-     * calling {@link Mqtt5Client#invokePublishAcknowledgement(Mqtt5PublishAcknowledgementControlHandle)}.</p>
-     *
-     * @return {@code true} if the user acquired manual control of the PUBACK, {@code false} otherwise.
-     */
-    private synchronized boolean wasControlAcquired() {
-        return controlAcquired;
-    }
-
-    /**
-     * Called by native/JNI code after the {@link Mqtt5ClientOptions.PublishEvents#onMessageReceived}
-     * callback returns to prevent post-callback use of
-     * {@link #acquirePublishAcknowledgementControl()}.
-     *
-     * <p>Zeroes out {@code controlId} so that any saved reference to this {@code PublishReturn}
-     * cannot call {@code acquirePublishAcknowledgementControl()} after the callback has returned.</p>
-     */
-    private synchronized void invalidateAfterCallback() {
-        controlId = 0;
-    }
-
     /**
      * This is only called in JNI to make a new PublishReturn with a PUBLISH packet.
      * The controlId is eagerly acquired by native code prior to
@@ -119,7 +82,6 @@ private synchronized void invalidateAfterCallback() {
     private PublishReturn(PublishPacket newPublishPacket, long controlId) {
         this.publishPacket = newPublishPacket;
         this.controlId = controlId;
-        this.controlAcquired = false;
         this.threadID = Thread.currentThread().getId();
     }
 }

src/native/java_class_ids.c

@@ -2265,20 +2265,17 @@ static void s_cache_mqtt5_publish_return(JNIEnv *env) {
         "(Lsoftware/amazon/awssdk/crt/mqtt5/packets/PublishPacket;J)V");
     AWS_FATAL_ASSERT(mqtt5_publish_return_properties.return_constructor_id);
     /*
-     * wasControlAcquired() called by native code after onMessageReceived returns to determine
-     * whether the user called acquirePublishAcknowledgementControl() during the callback.
-     * Returns true if the user took manual control of the PUBACK, false otherwise.
+     * acquirePublishAcknowledgementControl() called by native code after onMessageReceived returns
+     * to check whether the user took manual control during the callback. Returns a non-null handle
+     * if control was not yet acquired (native then auto-invokes the PUBACK), or null if the user
+     * already called it during the callback (user is responsible for invoking the PUBACK).
      */
-    mqtt5_publish_return_properties.return_was_control_acquired_id =
-        (*env)->GetMethodID(env, mqtt5_publish_return_properties.return_class, "wasControlAcquired", "()Z");
-    AWS_FATAL_ASSERT(mqtt5_publish_return_properties.return_was_control_acquired_id);
-    /*
-     * invalidateAfterCallback() called by native code after onMessageReceived returns to zero
-     * out controlId, preventing post-callback calls to acquirePublishAcknowledgementControl().
-     */
-    mqtt5_publish_return_properties.return_invalidate_after_callback_id =
-        (*env)->GetMethodID(env, mqtt5_publish_return_properties.return_class, "invalidateAfterCallback", "()V");
-    AWS_FATAL_ASSERT(mqtt5_publish_return_properties.return_invalidate_after_callback_id);
+    mqtt5_publish_return_properties.return_acquire_publish_acknowledgement_control_id = (*env)->GetMethodID(
+        env,
+        mqtt5_publish_return_properties.return_class,
+        "acquirePublishAcknowledgementControl",
+        "()Lsoftware/amazon/awssdk/crt/mqtt5/Mqtt5PublishAcknowledgementControlHandle;");
+    AWS_FATAL_ASSERT(mqtt5_publish_return_properties.return_acquire_publish_acknowledgement_control_id);
 }
 
 struct java_aws_mqtt5_on_stopped_return_properties mqtt5_on_stopped_return_properties;

src/native/java_class_ids.h

@@ -928,9 +928,9 @@ extern struct java_aws_mqtt5_publish_result_properties mqtt5_publish_result_prop
 /* mqtt5.PublishReturn */
 struct java_aws_mqtt5_publish_return_properties {
     jclass return_class;
-    jmethodID return_constructor_id;               /* (PublishPacket, long) */
-    jmethodID return_was_control_acquired_id;      /* ()Z true if user called acquirePublishAcknowledgementControl() */
-    jmethodID return_invalidate_after_callback_id; /* ()V zeroes controlId to prevent post-callback use */
+    jmethodID return_constructor_id;                             /* (PublishPacket, long) */
+    jmethodID return_acquire_publish_acknowledgement_control_id; /* ()Lsoftware/.../Mqtt5PublishAcknowledgementControlHandle;
+                                                                  */
 };
 extern struct java_aws_mqtt5_publish_return_properties mqtt5_publish_return_properties;
 

src/native/mqtt5_client.c

@@ -618,23 +618,22 @@ static void s_aws_mqtt5_client_java_publish_received(
     }
 
     /*
-     * After the callback returns:
-     * 1. Call invalidateAfterCallback() to zero out controlId in the PublishReturn, preventing
-     *    any post-callback call to acquirePublishAcknowledgementControl() from succeeding.
-     * 2. Call wasControlAcquired() to check whether the user took manual control during the
-     *    callback. If they did NOT, auto-invoke the PUBACK to avoid losing it.
+     * After the callback returns, call acquirePublishAcknowledgementControl() on the PublishReturn.
+     * - If the user already called it during the callback, controlId was zeroed out and the
+     *   method returns null. The user is responsible for invoking the PUBACK.
+     * - If the user did NOT call it, controlId is still non-zero and we get a non-null handle
+     *   back. We then auto-invoke the PUBACK to avoid losing it.
      */
     if (control_id != 0 && publish_packet_return_data != NULL) {
-        /* Invalidate the PublishReturn to prevent post-callback use */
-        (*env)->CallVoidMethod(
-            env, publish_packet_return_data, mqtt5_publish_return_properties.return_invalidate_after_callback_id);
-
-        /* Check whether the user called acquirePublishAcknowledgementControl() during the callback */
-        jboolean user_acquired_control = (*env)->CallBooleanMethod(
-            env, publish_packet_return_data, mqtt5_publish_return_properties.return_was_control_acquired_id);
+        jobject handle = (*env)->CallObjectMethod(
+            env,
+            publish_packet_return_data,
+            mqtt5_publish_return_properties.return_acquire_publish_acknowledgement_control_id);
+        aws_jni_check_and_clear_exception(env);
 
-        if (!user_acquired_control) {
-            /* User did NOT call acquirePublishAcknowledgementControl(); auto-invoke the PUBACK */
+        if (handle != NULL) {
+            /* User did NOT call acquirePublishAcknowledgementControl() during the callback;
+             * auto-invoke the PUBACK. */
             aws_mqtt5_client_invoke_publish_acknowledgement(java_client->client, control_id, NULL);
         }
     }