pybricks.tools: Make bluetooth tasks awaitable.

This is only used to await tasks that do not have to be cancelled,
like setting the remote light. Tasks with cancellation on timeout
(like connecting to the remote) remain blocking.

bricks/_common/sources.mk

@@ -106,7 +106,6 @@ PYBRICKS_PYBRICKS_SRC_C = $(addprefix pybricks/,\
 	util_pb/pb_conversions.c \
 	util_pb/pb_error.c \
 	util_pb/pb_serial_ev3dev.c \
-	util_pb/pb_task.c \
 	)
 
 # Pybricks I/O library

pybricks/common/pb_type_ble.c

@@ -17,8 +17,8 @@
 #include "py/runtime.h"
 
 #include <pybricks/common.h>
+#include <pybricks/tools.h>
 #include <pybricks/util_pb/pb_error.h>
-#include <pybricks/util_pb/pb_task.h>
 
 // The code currently passes integers and floats directly as bytes so requires
 // little-endian to get the correct ordering over the air.
@@ -47,6 +47,7 @@ static uint8_t num_observed_data;
 typedef struct {
     mp_obj_base_t base;
     uint8_t broadcast_channel;
+    pbio_task_t broadcast_task;
     observed_data_t observed_data[];
 } pb_obj_BLE_t;
 
@@ -276,12 +277,8 @@ STATIC mp_obj_t pb_module_ble_broadcast(size_t n_args, const mp_obj_t *args) {
     pbio_set_uint16_le(&value.v.data[2], LEGO_CID);
     value.v.data[4] = self->broadcast_channel;
 
-    pbio_task_t task;
-    pbdrv_bluetooth_start_broadcasting(&task, &value.v);
-
-    pb_wait_task(&task, -1);
-
-    return mp_const_none;
+    pbdrv_bluetooth_start_broadcasting(&self->broadcast_task, &value.v);
+    return pb_module_tools_pbio_task_wait_or_await(&self->broadcast_task);
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_VAR(pb_module_ble_broadcast_obj, 1, pb_module_ble_broadcast);
 
@@ -517,7 +514,7 @@ mp_obj_t pb_type_BLE_new(mp_obj_t broadcast_channel_in, mp_obj_t observe_channel
     if (num_channels > 0) {
         pbio_task_t task;
         pbdrv_bluetooth_start_observing(&task, handle_observe_event);
-        pb_wait_task(&task, -1);
+        pb_module_tools_pbio_task_do_blocking(&task, -1);
     }
 
     return MP_OBJ_FROM_PTR(self);

pybricks/iodevices/pb_type_iodevices_lwp3device.c

@@ -17,10 +17,10 @@
 
 #include <pybricks/common.h>
 #include <pybricks/parameters.h>
+#include <pybricks/tools.h>
 #include <pybricks/util_mp/pb_kwarg_helper.h>
 #include <pybricks/util_mp/pb_obj_helper.h>
 #include <pybricks/util_pb/pb_error.h>
-#include <pybricks/util_pb/pb_task.h>
 
 #include "py/mphal.h"
 #include "py/runtime.h"
@@ -82,7 +82,7 @@ STATIC void lwp3device_connect(const uint8_t hub_kind, const char *name, mp_int_
 
     pbdrv_bluetooth_set_notification_handler(handle_notification);
     pbdrv_bluetooth_scan_and_connect(&lwp3device->task, &lwp3device->context);
-    pb_wait_task(&lwp3device->task, timeout);
+    pb_module_tools_pbio_task_do_blocking(&lwp3device->task, timeout);
 }
 
 STATIC void lwp3device_assert_connected(void) {
@@ -144,7 +144,7 @@ STATIC mp_obj_t lwp3device_name(size_t n_args, const mp_obj_t *args) {
 
         // NB: operation is not cancelable, so timeout is not used
         pbdrv_bluetooth_write_remote(&lwp3device->task, &msg.value);
-        pb_wait_task(&lwp3device->task, -1);
+        pb_module_tools_pbio_task_do_blocking(&lwp3device->task, -1);
 
         // assuming write was successful instead of reading back from the handset
         memcpy(lwp3device->context.name, name, len);
@@ -181,9 +181,7 @@ STATIC mp_obj_t lwp3device_write(mp_obj_t self_in, mp_obj_t buf_in) {
     memcpy(msg.payload, bufinfo.buf, bufinfo.len);
 
     pbdrv_bluetooth_write_remote(&lwp3device->task, &msg.value);
-    pb_wait_task(&lwp3device->task, -1);
-
-    return MP_OBJ_NEW_SMALL_INT(bufinfo.len);
+    return pb_module_tools_pbio_task_wait_or_await(&lwp3device->task);
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_2(lwp3device_write_obj, lwp3device_write);
 

pybricks/pupdevices/pb_type_pupdevices_remote.c

@@ -18,10 +18,10 @@
 
 #include <pybricks/common.h>
 #include <pybricks/parameters.h>
+#include <pybricks/tools.h>
 #include <pybricks/util_mp/pb_kwarg_helper.h>
 #include <pybricks/util_mp/pb_obj_helper.h>
 #include <pybricks/util_pb/pb_error.h>
-#include <pybricks/util_pb/pb_task.h>
 
 #include "py/mphal.h"
 #include "py/runtime.h"
@@ -121,8 +121,7 @@ STATIC mp_obj_t pb_type_pupdevices_Remote_light_on(void *context, const pbio_col
     msg.payload[2] = msg.payload[2] * 3 / 8;
 
     pbdrv_bluetooth_write_remote(&remote->task, &msg.value);
-    pb_wait_task(&remote->task, -1);
-    return mp_const_none;
+    return pb_module_tools_pbio_task_wait_or_await(&remote->task);
 }
 
 STATIC void remote_connect(const char *name, mp_int_t timeout) {
@@ -150,7 +149,7 @@ STATIC void remote_connect(const char *name, mp_int_t timeout) {
 
     pbdrv_bluetooth_set_notification_handler(handle_notification);
     pbdrv_bluetooth_scan_and_connect(&remote->task, &remote->context);
-    pb_wait_task(&remote->task, timeout);
+    pb_module_tools_pbio_task_do_blocking(&remote->task, timeout);
 
     nlr_buf_t nlr;
     if (nlr_push(&nlr) == 0) {
@@ -177,21 +176,21 @@ STATIC void remote_connect(const char *name, mp_int_t timeout) {
         // set mode for left buttons
 
         pbdrv_bluetooth_write_remote(&remote->task, &msg.value);
-        pb_wait_task(&remote->task, -1);
+        pb_module_tools_pbio_task_do_blocking(&remote->task, -1);
 
         // set mode for right buttons
 
         msg.port = REMOTE_PORT_RIGHT_BUTTONS;
         pbdrv_bluetooth_write_remote(&remote->task, &msg.value);
-        pb_wait_task(&remote->task, -1);
+        pb_module_tools_pbio_task_do_blocking(&remote->task, -1);
 
         // set status light to RGB mode
 
         msg.port = REMOTE_PORT_STATUS_LIGHT;
         msg.mode = STATUS_LIGHT_MODE_RGB_0;
         msg.enable_notifications = 0;
         pbdrv_bluetooth_write_remote(&remote->task, &msg.value);
-        pb_wait_task(&remote->task, -1);
+        pb_module_tools_pbio_task_do_blocking(&remote->task, -1);
 
         // REVISIT: Could possibly use system color here to make remote match
         // hub status light. For now, the system color is hard-coded to blue.
@@ -310,7 +309,7 @@ STATIC mp_obj_t remote_name(size_t n_args, const mp_obj_t *args) {
 
         // NB: operation is not cancelable, so timeout is not used
         pbdrv_bluetooth_write_remote(&remote->task, &msg.value);
-        pb_wait_task(&remote->task, -1);
+        pb_module_tools_pbio_task_do_blocking(&remote->task, -1);
 
         // assuming write was successful instead of reading back from the handset
         memcpy(remote->context.name, name, len);

pybricks/tools.h

@@ -10,12 +10,18 @@
 
 #include "py/obj.h"
 
+#include <pbio/task.h>
+
 void pb_module_tools_init(void);
 
 bool pb_module_tools_run_loop_is_active(void);
 
 void pb_module_tools_assert_blocking(void);
 
+void pb_module_tools_pbio_task_do_blocking(pbio_task_t *task, mp_int_t timeout);
+
+mp_obj_t pb_module_tools_pbio_task_wait_or_await(pbio_task_t *task);
+
 extern const mp_obj_type_t pb_type_StopWatch;
 
 extern const mp_obj_type_t pb_type_Task;

pybricks/tools/pb_module_tools.c

@@ -13,6 +13,7 @@
 #include "py/stream.h"
 
 #include <pbio/int_math.h>
+#include <pbio/task.h>
 #include <pbsys/light.h>
 #include <pbsys/program_stop.h>
 
@@ -46,12 +47,6 @@ void pb_module_tools_assert_blocking(void) {
 // us share the same code with other awaitables. It also minimizes allocation.
 MP_REGISTER_ROOT_POINTER(mp_obj_t wait_awaitables);
 
-// Reset global state when user program starts.
-void pb_module_tools_init(void) {
-    MP_STATE_PORT(wait_awaitables) = mp_obj_new_list(0, NULL);
-    run_loop_is_active = false;
-}
-
 STATIC bool pb_module_tools_wait_test_completion(mp_obj_t obj, uint32_t end_time) {
     return mp_hal_ticks_ms() - end_time < UINT32_MAX / 2;
 }
@@ -86,6 +81,76 @@ STATIC mp_obj_t pb_module_tools_wait(size_t n_args, const mp_obj_t *pos_args, mp
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_KW(pb_module_tools_wait_obj, 0, pb_module_tools_wait);
 
+/**
+ * Waits for a task to complete.
+ *
+ * If an exception is raised while waiting, then the task is canceled.
+ *
+ * @param [in]  task    The task
+ * @param [in]  timeout The timeout in milliseconds or -1 to wait forever.
+ */
+void pb_module_tools_pbio_task_do_blocking(pbio_task_t *task, mp_int_t timeout) {
+
+    pb_module_tools_assert_blocking();
+
+    nlr_buf_t nlr;
+
+    if (nlr_push(&nlr) == 0) {
+        mp_uint_t start = mp_hal_ticks_ms();
+
+        while (timeout < 0 || mp_hal_ticks_ms() - start < (mp_uint_t)timeout) {
+            MICROPY_EVENT_POLL_HOOK
+
+            if (task->status != PBIO_ERROR_AGAIN) {
+                nlr_pop();
+                pb_assert(task->status);
+                return;
+            }
+        }
+
+        mp_raise_OSError(MP_ETIMEDOUT);
+        MP_UNREACHABLE
+    } else {
+        pbio_task_cancel(task);
+
+        while (task->status == PBIO_ERROR_AGAIN) {
+            MICROPY_VM_HOOK_LOOP
+        }
+
+        nlr_jump(nlr.ret_val);
+    }
+}
+
+// The awaitables associated with pbio tasks can originate from different
+// objects. At the moment, they are only associated with Bluetooth tasks, and
+// they cannot run at the same time. So we keep a single list of awaitables
+// here instead of with each Bluetooth-related MicroPython object.
+MP_REGISTER_ROOT_POINTER(mp_obj_t pbio_task_awaitables);
+
+STATIC bool pb_module_tools_pbio_task_test_completion(mp_obj_t obj, uint32_t end_time) {
+    pbio_task_t *task = MP_OBJ_TO_PTR(obj);
+
+    // Keep going if not done yet.
+    if (task->status == PBIO_ERROR_AGAIN) {
+        return false;
+    }
+
+    // If done, make sure it was successful.
+    pb_assert(task->status);
+    return true;
+}
+
+mp_obj_t pb_module_tools_pbio_task_wait_or_await(pbio_task_t *task) {
+    return pb_type_awaitable_await_or_wait(
+        MP_OBJ_FROM_PTR(task),
+        MP_STATE_PORT(pbio_task_awaitables),
+        pb_type_awaitable_end_time_none,
+        pb_module_tools_pbio_task_test_completion,
+        pb_type_awaitable_return_none,
+        pb_type_awaitable_cancel_none,
+        PB_TYPE_AWAITABLE_OPT_RAISE_ON_BUSY);
+}
+
 /**
  * Reads one byte from stdin without blocking.
  *
@@ -145,6 +210,13 @@ STATIC mp_obj_t pb_module_tools_run_task(size_t n_args, const mp_obj_t *pos_args
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_KW(pb_module_tools_run_task_obj, 1, pb_module_tools_run_task);
 
+// Reset global awaitable state when user program starts.
+void pb_module_tools_init(void) {
+    MP_STATE_PORT(wait_awaitables) = mp_obj_new_list(0, NULL);
+    MP_STATE_PORT(pbio_task_awaitables) = mp_obj_new_list(0, NULL);
+    run_loop_is_active = false;
+}
+
 #if PYBRICKS_PY_TOOLS_HUB_MENU
 
 STATIC void pb_module_tools_hub_menu_display_symbol(mp_obj_t symbol) {