pbio/sys/host: fix pbsys_host_stdout_write()

Fix issue with pbsys_host_stdout_write() not being cancellable. The
static state variables in the function required that the function be
called repeatedly until it returned something other than PBIO_ERROR_AGAIN.
However, in MicroPython, an exception may be raised in between calls,
which would leave the state variables in an inconsistent state causing
incorrect behavior on the next call.

We can save code size on platforms that only have one transport by only
calling the relevant function instead of trying to duplicate the stream
to a non-existent transport.

Then, to handle platforms with multiple transports, rework the logic to
be able to send data to multiple transports without having to maintain
state between calls. This is done by ensuring that we only send the
same amount of data to each transport.

Author: dlech

bricks/_common_stm32/mphalport.c

@@ -60,10 +60,22 @@ int mp_hal_stdin_rx_chr(void) {
 
 // Send string of given length
 mp_uint_t mp_hal_stdout_tx_strn(const char *str, size_t len) {
-    while (pbsys_host_tx((const uint8_t *)str, len) == PBIO_ERROR_AGAIN) {
+    size_t remaining = len;
+
+    while (remaining) {
+        uint32_t size = remaining;
+        pbio_error_t err = pbsys_host_stdout_write((const uint8_t *)str, &size);
+        if (err == PBIO_SUCCESS) {
+            str += size;
+            remaining -= size;
+        } else if (err != PBIO_ERROR_AGAIN) {
+            // Ignoring error for now. This means stdout is lost if Bluetooth/USB
+            // is disconnected.
+            return len - remaining;
+        }
+
         MICROPY_EVENT_POLL_HOOK
     }
-    // Not raising the error. This means stdout lost if host is not connected.
 
     return len;
 }

lib/pbio/drv/usb/usb_stm32.c

@@ -145,8 +145,8 @@ void pbdrv_usb_stm32_handle_vbus_irq(bool active) {
  * @param data  [in]        The data to be sent.
  * @param size  [in, out]   The size of @p data in bytes. After return, @p size
  *                          contains the number of bytes actually written.
- * @return                  ::PBIO_SUCCESS if all @p data was queued, ::PBIO_ERROR_AGAIN
- *                          if not all @p data could not be queued at this time (e.g. buffer
+ * @return                  ::PBIO_SUCCESS if some @p data was queued, ::PBIO_ERROR_AGAIN
+ *                          if no @p data could not be queued at this time (e.g. buffer
  *                          is full), ::PBIO_ERROR_INVALID_OP if there is not an
  *                          active USB connection or ::PBIO_ERROR_NOT_SUPPORTED
  *                          if this platform does not support USB.
@@ -158,7 +158,6 @@ pbio_error_t pbdrv_usb_stdout_tx(const uint8_t *data, uint32_t *size) {
 
     uint8_t *ptr = usb_stdout_buf;
     uint32_t ptr_len = sizeof(usb_stdout_buf);
-    uint32_t full_size = *size;
 
     // TODO: return PBIO_ERROR_INVALID_OP if app flag is not set. Also need a
     // timeout in case the app crashes and doesn't clear the flag on exit.
@@ -174,14 +173,24 @@ pbio_error_t pbdrv_usb_stdout_tx(const uint8_t *data, uint32_t *size) {
     *ptr++ = PBIO_PYBRICKS_EVENT_WRITE_STDOUT;
     ptr_len--;
 
-    *size = MIN(full_size, ptr_len);
+    *size = MIN(*size, ptr_len);
     memcpy(ptr, data, *size);
 
     usb_stdout_sz = 1 + 1 + *size;
 
     process_poll(&pbdrv_usb_process);
 
-    return *size == full_size ? PBIO_SUCCESS : PBIO_ERROR_AGAIN;
+    return PBIO_SUCCESS;
+}
+
+uint32_t pbdrv_usb_stdout_tx_available(void) {
+    // TODO: return UINT32_MAX if app flag is not set.
+
+    if (usb_stdout_sz) {
+        return 0;
+    }
+
+    return sizeof(usb_stdout_buf) - 2; // 2 bytes for header
 }
 
 /**

lib/pbio/include/pbdrv/usb.h

@@ -46,6 +46,16 @@ pbdrv_usb_bcd_t pbdrv_usb_get_bcd(void);
  */
 pbio_error_t pbdrv_usb_stdout_tx(const uint8_t *data, uint32_t *size);
 
+/**
+ * Gets the number of bytes that can be queued for sending stdout via USB.
+ *
+ * Returns UINT32_MAX if there is no USB connection or no app is subscribed to
+ * stdout.
+ *
+ * @return              The number of bytes that can be queued.
+ */
+uint32_t pbdrv_usb_stdout_tx_available(void);
+
 /**
  * Indicates if the USB stdout stream is idle.
  * @return              true if the USB stdout stream is idle.
@@ -62,6 +72,10 @@ static inline pbio_error_t pbdrv_usb_stdout_tx(const uint8_t *data, uint32_t *si
     return PBIO_SUCCESS;
 }
 
+static inline uint32_t pbdrv_usb_stdout_tx_available(void) {
+    return UINT32_MAX;
+}
+
 static inline bool pbdrv_usb_stdout_tx_is_idle(void) {
     return true;
 }

lib/pbio/include/pbsys/host.h

@@ -32,7 +32,7 @@ void pbsys_host_stdin_set_callback(pbsys_host_stdin_event_callback_t callback);
 void pbsys_host_stdin_flush(void);
 uint32_t pbsys_host_stdin_get_available(void);
 pbio_error_t pbsys_host_stdin_read(uint8_t *data, uint32_t *size);
-pbio_error_t pbsys_host_tx(const uint8_t *data, uint32_t size);
+pbio_error_t pbsys_host_stdout_write(const uint8_t *data, uint32_t *size);
 bool pbsys_host_tx_is_idle(void);
 
 #else // PBSYS_CONFIG_HOST
@@ -44,7 +44,7 @@ bool pbsys_host_tx_is_idle(void);
 #define pbsys_host_stdin_flush()
 #define pbsys_host_stdin_get_available() 0
 #define pbsys_host_stdin_read(data, size) PBIO_ERROR_NOT_SUPPORTED
-#define pbsys_host_tx(data, size) { (void)(data); (void)(size); PBIO_ERROR_NOT_SUPPORTED; }
+#define pbsys_host_stdout_write(data, size) { (void)(data); (void)(size); PBIO_ERROR_NOT_SUPPORTED; }
 #define pbsys_host_tx_is_idle() false
 
 #endif // PBSYS_CONFIG_HOST

lib/pbio/sys/bluetooth.c

@@ -67,17 +67,15 @@ void pbsys_bluetooth_init(void) {
  * @param data  [in]        The data to be sent.
  * @param size  [in, out]   The size of @p data in bytes. After return, @p size
  *                          contains the number of bytes actually written.
- * @return                  ::PBIO_SUCCESS if all @p data was queued,
- *                          ::PBIO_ERROR_AGAIN if not all @p data could not be
+ * @return                  ::PBIO_SUCCESS if some @p data was queued,
+ *                          ::PBIO_ERROR_AGAIN if no @p data could be
  *                          queued at this time (e.g. buffer is full),
  *                          ::PBIO_ERROR_INVALID_OP if there is not an
  *                          active Bluetooth connection or ::PBIO_ERROR_NOT_SUPPORTED
  *                          if this platform does not support Bluetooth.
  */
 pbio_error_t pbsys_bluetooth_tx(const uint8_t *data, uint32_t *size) {
 
-    uint32_t full_size = *size;
-
     // make sure we have a Bluetooth connection
     if (!pbdrv_bluetooth_is_connected(PBDRV_BLUETOOTH_CONNECTION_PYBRICKS)) {
         return PBIO_ERROR_INVALID_OP;
@@ -94,15 +92,30 @@ pbio_error_t pbsys_bluetooth_tx(const uint8_t *data, uint32_t *size) {
         stdout_msg.is_queued = true;
     }
 
-    if ((*size = lwrb_write(&stdout_ring_buf, data, full_size)) == 0) {
+    if ((*size = lwrb_write(&stdout_ring_buf, data, *size)) == 0) {
         return PBIO_ERROR_AGAIN;
     }
 
     // poke the process to start tx soon-ish. This way, we can accumulate up to
     // MAX_CHAR_SIZE bytes before actually transmitting
     pbsys_bluetooth_process_poll();
 
-    return *size == full_size ? PBIO_SUCCESS : PBIO_ERROR_AGAIN;
+    return PBIO_SUCCESS;
+}
+
+/**
+ * Gets the number of bytes that can be queued for transmission via Bluetooth.
+ *
+ * If there is no connection, this will return %UINT32_MAX.
+ *
+ * @returns The number of bytes that can be queued.
+ */
+uint32_t pbsys_bluetooth_tx_available(void) {
+    if (!pbdrv_bluetooth_is_connected(PBDRV_BLUETOOTH_CONNECTION_PYBRICKS)) {
+        return UINT32_MAX;
+    }
+
+    return lwrb_get_free(&stdout_ring_buf);
 }
 
 /**

lib/pbio/sys/bluetooth.h

@@ -16,6 +16,7 @@ void pbsys_bluetooth_process_poll(void);
 
 void pbsys_bluetooth_init(void);
 pbio_error_t pbsys_bluetooth_tx(const uint8_t *data, uint32_t *size);
+uint32_t pbsys_bluetooth_tx_available(void);
 bool pbsys_bluetooth_tx_is_idle(void);
 
 #else // PBSYS_CONFIG_BLUETOOTH
@@ -25,6 +26,9 @@ bool pbsys_bluetooth_tx_is_idle(void);
 static inline pbio_error_t pbsys_bluetooth_tx(const uint8_t *data, uint32_t *size) {
     return PBIO_ERROR_NOT_SUPPORTED;
 }
+static inline uint32_t pbsys_bluetooth_tx_available(void) {
+    return UINT32_MAX;
+}
 static inline bool pbsys_bluetooth_tx_is_idle(void) {
     return false;
 }

lib/pbio/sys/host.c

@@ -7,6 +7,7 @@
 
 #include <lwrb/lwrb.h>
 
+#include <pbdrv/bluetooth.h>
 #include <pbdrv/usb.h>
 #include <pbsys/host.h>
 
@@ -101,53 +102,57 @@ pbio_error_t pbsys_host_stdin_read(uint8_t *data, uint32_t *size) {
 }
 
 /**
- * Transmits data over Bluetooth and USB.
+ * Transmits data over any connected transport that is subscribed to Pybricks
+ * protocol events.
  *
- * Should be called in a loop with the same arguments until it no longer
- * returns ::PBIO_ERROR_AGAIN.
+ * This may perform partial writes. Callers should check the number of bytes
+ * actually written and call again with the remaining data until all data is
+ * written.
  *
- * @param data  [in]   The data to transmit.
- * @param size  [in]   The size of the data to transmit.
- * @return             ::PBIO_ERROR_AGAIN if the data is still being transmitted
- *                     ::PBIO_SUCCESS if complete or failed.
+ * @param data  [in]        The data to transmit.
+ * @param size  [inout]     The size of the data to transmit. Upon success, this
+ *                          contains the number of bytes actually processed.
+ * @return                  ::PBIO_ERROR_INVALID_OP if there is no active transport,
+ *                          ::PBIO_ERROR_AGAIN if no @p data could be queued,
+ *                          ::PBIO_SUCCESS if at least some data was queued.
  */
-pbio_error_t pbsys_host_tx(const uint8_t *data, uint32_t size) {
-
-    static bool transmitting = false;
-    static uint32_t tx_done_ble;
-    static uint32_t tx_done_usb;
-
-    if (!transmitting) {
-        tx_done_ble = 0;
-        tx_done_usb = 0;
-        transmitting = true;
+pbio_error_t pbsys_host_stdout_write(const uint8_t *data, uint32_t *size) {
+    #if PBSYS_CONFIG_BLUETOOTH && (!PBDRV_CONFIG_USB || PBDRV_CONFIG_USB_CHARGE_ONLY)
+    return pbsys_bluetooth_tx(data, size);
+    #elif !PBSYS_CONFIG_BLUETOOTH && PBDRV_CONFIG_USB && !PBDRV_CONFIG_USB_CHARGE_ONLY
+    return pbdrv_usb_stdout_tx(data, size);
+    #elif PBSYS_CONFIG_BLUETOOTH && PBDRV_CONFIG_USB && !PBDRV_CONFIG_USB_CHARGE_ONLY
+
+    uint32_t bt_avail = pbsys_bluetooth_tx_available();
+    uint32_t usb_avail = pbdrv_usb_stdout_tx_available();
+    uint32_t available = MIN(UINT32_MAX, MIN(bt_avail, usb_avail));
+
+    // If all tx_available() calls returned UINT32_MAX, then there is one listening.
+    if (available == UINT32_MAX) {
+        return PBIO_ERROR_INVALID_OP;
     }
-
-    pbio_error_t err_ble = PBIO_SUCCESS;
-    pbio_error_t err_usb = PBIO_SUCCESS;
-    uint32_t size_now;
-
-    if (tx_done_ble < size) {
-        size_now = size - tx_done_ble;
-        err_ble = pbsys_bluetooth_tx(data + tx_done_ble, &size_now);
-        tx_done_ble += size_now;
+    // If one or more tx_available() calls returned 0, then we need to wait.
+    if (available == 0) {
+        return PBIO_ERROR_AGAIN;
     }
 
-    if (tx_done_usb < size) {
-        size_now = size - tx_done_usb;
-        err_usb = pbdrv_usb_stdout_tx(data + tx_done_usb, &size_now);
-        tx_done_usb += size_now;
-    }
+    // Limit size to smallest available space from all transports so that we
+    // don't do partial writes to one transport and not the other.
+    *size = MIN(*size, available);
 
-    // Keep waiting as long as at least has not completed or errored.
-    if (err_ble == PBIO_ERROR_AGAIN || err_usb == PBIO_ERROR_AGAIN) {
-        return PBIO_ERROR_AGAIN;
-    }
+    // Unless something became disconnected in an interrupt handler, these
+    // functions should always succeed since we already checked tx_available().
+    // And if both somehow got disconnected at the same time, it is not a big deal
+    // if we return PBIO_SUCCESS without actually sending anything.
+    (void)pbsys_bluetooth_tx(data, size);
+    (void)pbdrv_usb_stdout_tx(data, size);
+
+    return PBIO_SUCCESS;
 
-    // Both of them are either complete or failed. The caller of this function
-    // does not currently raise errors, so we just return success.
-    transmitting = false;
+    #else
+    // stdout goes to /dev/null
     return PBIO_SUCCESS;
+    #endif
 }
 
 bool pbsys_host_tx_is_idle(void) {