uringmachine 0.28.3 → 0.29.0

data/ext/um/um.h

@@ -77,16 +77,24 @@ enum um_op_kind {
   OP_TIMEOUT_MULTISHOT,
 };
 
+enum um_stream_mode {
+  STREAM_BP_READ,
+  STREAM_BP_RECV,
+  STREAM_SSL,
+  STREAM_STRING,
+  STREAM_IO_BUFFER
+};
 
 #define OP_F_CQE_SEEN       (1U << 0) // CQE has been seen
 #define OP_F_CQE_DONE       (1U << 1) // CQE has been seen and operation is done
 #define OP_F_SCHEDULED      (1U << 2) // op is on runqueue
 #define OP_F_CANCELED       (1U << 3) // op is cancelled (disregard CQE results)
 #define OP_F_MULTISHOT      (1U << 4) // op is multishot
-#define OP_F_ASYNC          (1U << 5) // op is async (no fiber is scheduled to be resumed on completion)
-#define OP_F_TRANSIENT      (1U << 6) // op is on transient list (for GC purposes)
+#define OP_F_ASYNC          (1U << 5) // op is async (don't schedule fiber)
+#define OP_F_TRANSIENT      (1U << 6) // op is on transient list (for GC marking)
 #define OP_F_FREE_IOVECS    (1U << 7) // op->iovecs should be freed on release
 #define OP_F_SKIP           (1U << 8) // op should be skipped when pulled from runqueue
+#define OP_F_BUFFER_POOL    (1U << 9) // multishot op using buffer pool
 
 #define OP_F_SELECT_POLLIN  (1U << 13) // select POLLIN
 #define OP_F_SELECT_POLLOUT (1U << 14) // select POLLOUT
@@ -101,9 +109,36 @@ enum um_op_kind {
 #define OP_TRANSIENT_P(op)  ((op)->flags & OP_F_TRANSIENT)
 #define OP_SKIP_P(op)       ((op)->flags & OP_F_SKIP)
 
+#define BP_BGID                     0xF00B
+#define BP_BR_ENTRIES               1024
+
+#define BP_INITIAL_BUFFER_SIZE      (1U << 14) // 16KB
+#define BP_INITIAL_COMMIT_THRESHOLD (1U << 16) // 64KB
+#define BP_MAX_BUFFER_SIZE          (1U << 20) // 1MB
+
+#define BP_MAX_COMMIT_THRESHOLD     (BP_MAX_BUFFER_SIZE * BP_BR_ENTRIES) // 1GB
+#define BP_AVAIL_BID_BITMAP_WORDS   (BP_BR_ENTRIES / 64)
+
+struct um_buffer {
+  size_t len;
+  size_t pos;
+  uint ref_count;
+  struct um_buffer *next;
+  char buf[];
+};
+
+struct um_segment {
+  char *ptr;
+  size_t len;
+
+  struct um_buffer *buffer;
+  struct um_segment *next;
+};
+
 struct um_op_result {
   __s32 res;
   __u32 flags;
+  struct um_segment *segment;
   struct um_op_result *next;
 };
 
@@ -128,15 +163,10 @@ struct um_op {
     struct iovec *iovecs; // used for vectorized write/send
     siginfo_t siginfo; // used for waitid
     int int_value; // used for getsockopt
+    size_t bp_commit_threshold; // buffer pool commit threshold
   };
 };
 
-struct um_buffer {
-  struct um_buffer *next;
-  void *ptr;
-  long len;
-};
-
 struct buf_ring_descriptor {
   struct io_uring_buf_ring *br;
   size_t br_size;
@@ -147,19 +177,25 @@ struct buf_ring_descriptor {
 };
 
 struct um_metrics {
-  ulong total_ops;        // total ops submitted
-  ulong total_switches;   // total fiber switches
-  ulong total_waits;      // total number of CQE waits
-
-  uint ops_pending;       // number of pending ops
-  uint ops_unsubmitted;   // number of unsubmitted
-  uint ops_runqueue;      // number of ops in runqueue
-  uint ops_free;          // number of ops in freelist
-  uint ops_transient;     // number of ops in transient list
-
-  double time_total_wait; // total CPU time waiting for CQEs
-  double time_last_cpu; // last seen time stamp
-  double time_first_cpu; // last seen time stamp
+  ulong total_ops;                // total ops submitted
+  ulong total_switches;           // total fiber switches
+  ulong total_waits;              // total number of CQE waits
+
+  uint ops_pending;               // number of pending ops
+  uint ops_unsubmitted;           // number of unsubmitted
+  uint ops_runqueue;              // number of ops in runqueue
+  uint ops_free;                  // number of ops in freelist
+  uint ops_transient;             // number of ops in transient list
+
+  uint buffers_allocated;         // number of allocated buffers
+  uint buffers_free;              // number of available buffers
+  uint segments_free;             // free segments
+  size_t buffer_space_allocated;  // total allocated buffer space
+  size_t buffer_space_commited;   // commited buffer space
+
+  double time_total_wait;         // total CPU time waiting for CQEs
+  double time_last_cpu;           // last seen time stamp
+  double time_first_cpu;          // last seen time stamp
 };
 
 #define BUFFER_RING_MAX_COUNT 10
@@ -167,8 +203,6 @@ struct um_metrics {
 struct um {
   VALUE self;
 
-  struct um_buffer *buffer_freelist;
-
   struct io_uring ring;
 
   uint ring_initialized; // is the ring initialized successfully
@@ -196,6 +230,17 @@ struct um {
 
   struct um_op *op_freelist;
   struct um_op_result *result_freelist;
+  struct um_segment *segment_freelist;
+
+  struct io_uring_buf_ring *bp_br;
+  size_t bp_buffer_size;
+  size_t bp_commit_threshold;
+  struct um_buffer **bp_commited_buffers;
+  uint64_t bp_avail_bid_bitmap[BP_AVAIL_BID_BITMAP_WORDS];
+
+  struct um_buffer *bp_buffer_freelist;
+  uint bp_buffer_count;
+  size_t bp_total_commited;
 };
 
 struct um_mutex {
@@ -227,11 +272,21 @@ struct um_async_op {
 };
 
 struct um_stream {
+  VALUE self;
   struct um *machine;
-  int fd;
-  VALUE buffer;
-  ulong len;
-  ulong pos;
+
+  enum um_stream_mode mode;
+  union {
+    int fd;
+    VALUE target;
+  };
+
+  struct um_buffer *working_buffer;
+  struct um_op *op;
+  struct um_segment *head;
+  struct um_segment *tail;
+  size_t pending_len;
+  size_t pos;
   int eof;
 };
 
@@ -271,12 +326,15 @@ void um_op_list_compact(struct um *machine, struct um_op *head);
 void um_op_multishot_results_push(struct um *machine, struct um_op *op, __s32 res, __u32 flags);
 void um_op_multishot_results_clear(struct um *machine, struct um_op *op);
 
+struct um_segment *um_segment_alloc(struct um *machine);
+void um_segment_free(struct um *machine, struct um_segment *segment);
+
 void um_runqueue_push(struct um *machine, struct um_op *op);
 struct um_op *um_runqueue_shift(struct um *machine);
 
 struct um_buffer *um_buffer_checkout(struct um *machine, int len);
-void um_buffer_checkin(struct um *machine, struct um_buffer *buffer);
-void um_free_buffer_linked_list(struct um *machine);
+void bp_buffer_checkin(struct um *machine, struct um_buffer *buffer);
+void bp_discard_buffer_freelist(struct um *machine);
 
 struct __kernel_timespec um_double_to_timespec(double value);
 double um_timestamp_to_double(__s64 tv_sec, __u32 tv_nsec);
@@ -375,8 +433,10 @@ VALUE um_queue_pop(struct um *machine, struct um_queue *queue);
 VALUE um_queue_unshift(struct um *machine, struct um_queue *queue, VALUE value);
 VALUE um_queue_shift(struct um *machine, struct um_queue *queue);
 
-VALUE stream_get_line(struct um_stream *stream, VALUE buf, ssize_t maxlen);
-VALUE stream_get_string(struct um_stream *stream, VALUE buf, ssize_t len);
+void stream_teardown(struct um_stream *stream);
+void stream_clear(struct um_stream *stream);
+VALUE stream_get_line(struct um_stream *stream, VALUE buf, size_t maxlen);
+VALUE stream_get_string(struct um_stream *stream, VALUE out_buffer, ssize_t len, size_t inc, int safe_inc);
 VALUE stream_skip(struct um_stream *stream, size_t len);
 VALUE resp_decode(struct um_stream *stream, VALUE out_buffer);
 void resp_encode(struct um_write_buffer *buf, VALUE obj);
@@ -396,6 +456,18 @@ void um_sidecar_signal_wake(struct um *machine);
 
 void um_ssl_set_bio(struct um *machine, VALUE ssl_obj);
 int um_ssl_read(struct um *machine, VALUE ssl, VALUE buf, int maxlen);
+int um_ssl_read_raw(struct um *machine, VALUE ssl_obj, char *ptr, int maxlen);
 int um_ssl_write(struct um *machine, VALUE ssl, VALUE buf, int len);
 
+void bp_setup(struct um *machine);
+void bp_teardown(struct um *machine);
+
+struct um_segment *bp_get_op_result_segment(struct um *machine, struct um_op *op, __s32 res, __u32 flags);
+void um_segment_checkin(struct um *machine, struct um_segment *segment);
+void bp_handle_enobufs(struct um *machine);
+void bp_ensure_commit_level(struct um *machine);
+struct um_buffer *bp_buffer_checkout(struct um *machine);
+void bp_buffer_checkin(struct um *machine, struct um_buffer *buffer);
+struct um_segment *bp_buffer_consume(struct um *machine, struct um_buffer *buffer, size_t len);
+
 #endif // UM_H

data/ext/um/um_buffer_pool.c

@@ -0,0 +1,246 @@
+#include <string.h>
+#include "um.h"
+
+#define UM_SEGMENT_ALLOC_BATCH_SIZE   256
+
+inline struct um_buffer *buffer_alloc(struct um *machine) {
+  struct um_buffer *buffer = malloc(sizeof(struct um_buffer) + machine->bp_buffer_size);
+  if (unlikely(!buffer)) {
+    fprintf(stderr, "!ENOMEM!\n");
+    exit(1);
+    rb_syserr_fail(errno, strerror(errno));
+  }
+
+  machine->metrics.buffers_allocated += 1;
+  machine->metrics.buffer_space_allocated += machine->bp_buffer_size;
+  buffer->len = machine->bp_buffer_size;
+  buffer->ref_count = 0;
+  return buffer;
+}
+
+inline struct um_buffer *bp_buffer_checkout(struct um *machine) {
+  struct um_buffer *buffer = machine->bp_buffer_freelist;
+  if (buffer) {
+    struct um_buffer *next = buffer->next;
+    machine->bp_buffer_freelist = next;
+    machine->metrics.buffers_free--;
+  }
+  else
+    buffer = buffer_alloc(machine);
+
+  buffer->ref_count++;
+  buffer->pos = 0;
+  buffer->next = NULL;
+
+  return buffer;
+}
+
+inline void buffer_free(struct um *machine, struct um_buffer *buffer) {
+  if (unlikely(!machine->ring_initialized || buffer->len != machine->bp_buffer_size)) {
+    // The machine is being shut down or working buffer size has changed, so the
+    // buffer can be freed.
+    machine->metrics.buffers_allocated -= 1;
+    machine->metrics.buffer_space_allocated -= buffer->len;
+    free(buffer);
+  }
+  else {
+    // otherwise, keep it on the freelist
+    buffer->next = machine->bp_buffer_freelist;
+    machine->bp_buffer_freelist = buffer;
+    machine->metrics.buffers_free++;
+  }
+}
+
+inline void bp_buffer_checkin(struct um *machine, struct um_buffer *buffer) {
+  assert(buffer->ref_count > 0);
+  buffer->ref_count--;
+  if (!buffer->ref_count) buffer_free(machine, buffer);
+}
+
+inline void bp_discard_buffer_freelist(struct um *machine) {
+  while (machine->bp_buffer_freelist) {
+    struct um_buffer *buffer = machine->bp_buffer_freelist;
+    struct um_buffer *next = buffer->next;
+
+    machine->metrics.buffers_allocated -= 1;
+    machine->metrics.buffer_space_allocated -= buffer->len;
+
+    free(buffer);
+    machine->bp_buffer_freelist = next;
+  }
+}
+
+inline void bp_setup(struct um *machine) {
+  int ret;
+  machine->bp_br = io_uring_setup_buf_ring(&machine->ring, BP_BR_ENTRIES, BP_BGID, IOU_PBUF_RING_INC, &ret);
+  if (unlikely(!machine->bp_br)) rb_syserr_fail(ret, strerror(ret));
+
+  machine->bp_buffer_size = BP_INITIAL_BUFFER_SIZE;
+  machine->bp_commit_threshold = BP_INITIAL_COMMIT_THRESHOLD;
+  machine->bp_commited_buffers = malloc(sizeof(struct um_buffer) * BP_BR_ENTRIES);
+  memset(machine->bp_commited_buffers, 0, sizeof(struct um_buffer) * BP_BR_ENTRIES);
+  memset(machine->bp_avail_bid_bitmap, 0xFF, sizeof(machine->bp_avail_bid_bitmap));
+
+  machine->bp_buffer_freelist = NULL;
+  machine->bp_buffer_count = 0;
+  machine->bp_total_commited = 0;
+}
+
+inline void bp_teardown(struct um *machine) {
+  bp_discard_buffer_freelist(machine);
+  for (int i = 0; i < BP_BR_ENTRIES; i++) {
+    struct um_buffer *buffer = machine->bp_commited_buffers[i];
+    if (buffer) bp_buffer_checkin(machine, buffer);
+  }
+  free(machine->bp_commited_buffers);
+}
+
+inline int bitmap_find_first_set_bit(uint64_t *bitmap, size_t word_count) {
+  uint word = 0;
+  for (; word < word_count && !bitmap[word]; word++)
+  if (word == word_count) return -1;
+
+  int bit = ffsll(bitmap[word]);
+  return (word * 64 + bit - 1);
+}
+
+inline void bitmap_set(uint64_t *bitmap, int idx) {
+  uint word = idx / 64;
+  uint bit = idx % 64;
+  bitmap[word] |= (1U << bit);
+}
+
+inline void bitmap_unset(uint64_t *bitmap, int idx) {
+  uint word = idx / 64;
+  uint bit = idx % 64;
+  bitmap[word] &= ~(1U << bit);
+}
+
+inline int get_available_bid(struct um *machine) {
+  return bitmap_find_first_set_bit(machine->bp_avail_bid_bitmap, BP_AVAIL_BID_BITMAP_WORDS);
+}
+
+// Finds an available bid, checks out a buffer and commits it for kernel usage.
+// Returns true if successful.
+static inline int commit_buffer(struct um *machine, int added) {
+  static int buf_mask;
+  if (!buf_mask) buf_mask = io_uring_buf_ring_mask(BP_BR_ENTRIES);
+
+  int bid = get_available_bid(machine);
+  if (bid < 0) return false;
+
+  bitmap_unset(machine->bp_avail_bid_bitmap, bid);
+  struct um_buffer *buffer = bp_buffer_checkout(machine);
+  assert(buffer->ref_count == 0);
+  buffer->ref_count = 1;
+
+  machine->bp_buffer_count++;
+  machine->bp_commited_buffers[bid] = buffer;
+  machine->bp_total_commited += buffer->len;
+  machine->metrics.buffer_space_commited += buffer->len;
+
+  io_uring_buf_ring_add(machine->bp_br, buffer->buf, buffer->len, bid, buf_mask, added);
+  return true;
+}
+
+// Removes buffer from bid slot, the buffer having already been entirely
+// consumed by the kernel.
+inline void uncommit_buffer(struct um *machine, struct um_op *op, struct um_buffer *buffer, int bid) {
+  machine->bp_buffer_count--;
+  machine->bp_commited_buffers[bid] = NULL;
+  bitmap_set(machine->bp_avail_bid_bitmap, bid);
+  bp_buffer_checkin(machine, buffer);
+}
+
+inline struct um_buffer *get_buffer(struct um *machine, int bid) {
+  struct um_buffer *buffer = machine->bp_commited_buffers[bid];
+  assert(buffer);
+  return buffer;
+}
+
+inline int should_commit_more_p(struct um *machine) {
+  return (machine->bp_buffer_count < BP_BR_ENTRIES) &&
+         (machine->bp_total_commited < machine->bp_commit_threshold);
+}
+
+inline void bp_ensure_commit_level(struct um *machine) {
+  int added = 0;
+  while (should_commit_more_p(machine)) {
+    if (likely(commit_buffer(machine, added))) added++;
+  }
+  if (added) io_uring_buf_ring_advance(machine->bp_br, added);
+
+  // if we get to this point, there's nothing more we can do because we used up
+  // all buffer ring entries. We need to wait for the kernel to consume buffers
+  // in order to put in more. When a ENOBUFS error is received,
+  // bp_handle_enobufs is called and among other things increases the buffer
+  // size.
+}
+
+inline void bp_handle_enobufs(struct um *machine) {
+  if (unlikely(machine->bp_commit_threshold >= BP_MAX_COMMIT_THRESHOLD))
+    rb_raise(eUMError, "Buffer starvation");
+
+  machine->bp_commit_threshold *= 2;
+  while (machine->bp_buffer_size < machine->bp_commit_threshold / 4)
+    machine->bp_buffer_size *= 2;
+  bp_discard_buffer_freelist(machine);
+}
+
+inline struct um_segment *um_segment_alloc(struct um *machine) {
+  if (machine->segment_freelist) {
+    struct um_segment *segment = machine->segment_freelist;
+    machine->segment_freelist = segment->next;
+    machine->metrics.segments_free--;
+    segment->next = NULL;
+    return segment;
+  }
+
+  struct um_segment *batch = malloc(sizeof(struct um_segment) * UM_SEGMENT_ALLOC_BATCH_SIZE);
+  memset(batch, 0, sizeof(struct um_segment) * UM_SEGMENT_ALLOC_BATCH_SIZE);
+  for (int i = 1; i < (UM_SEGMENT_ALLOC_BATCH_SIZE - 1); i++) {
+    batch[i].next = &batch[i + 1];
+  }
+  machine->segment_freelist = batch + 1;
+  machine->metrics.segments_free += (UM_SEGMENT_ALLOC_BATCH_SIZE - 1);
+  return batch;
+}
+
+inline void um_segment_free(struct um *machine, struct um_segment *segment) {
+  segment->next = machine->segment_freelist;
+  machine->segment_freelist = segment;
+  machine->metrics.segments_free++;
+}
+
+inline struct um_segment *bp_buffer_consume(struct um *machine, struct um_buffer *buffer, size_t len) {
+  struct um_segment *segment = um_segment_alloc(machine);
+  segment->ptr = buffer->buf + buffer->pos;
+  segment->len = len;
+  segment->buffer = buffer;
+  buffer->pos += len;
+  buffer->ref_count++;
+  return segment;
+}
+
+struct um_segment *bp_get_op_result_segment(struct um *machine, struct um_op *op, __s32 res, __u32 flags) {
+  assert(res >= 0);
+  if (!res) return NULL;
+
+  uint bid = flags >> IORING_CQE_BUFFER_SHIFT;
+  struct um_buffer *buffer = get_buffer(machine, bid);
+
+  struct um_segment *segment = bp_buffer_consume(machine, buffer, res);
+  machine->bp_total_commited -= res;
+  machine->metrics.buffer_space_commited -= res;
+  bp_ensure_commit_level(machine);
+
+  if (!(flags & IORING_CQE_F_BUF_MORE))
+    uncommit_buffer(machine, op, buffer, bid);
+
+  return segment;
+}
+
+inline void um_segment_checkin(struct um *machine, struct um_segment *segment) {
+  bp_buffer_checkin(machine, segment->buffer);
+  um_segment_free(machine, segment);
+}

data/ext/um/um_class.c

@@ -22,6 +22,12 @@ VALUE SYM_ops_free;
 VALUE SYM_ops_transient;
 VALUE SYM_time_total_cpu;
 VALUE SYM_time_total_wait;
+VALUE SYM_buffer_groups;
+VALUE SYM_buffers_allocated;
+VALUE SYM_buffers_free;
+VALUE SYM_segments_free;
+VALUE SYM_buffer_space_allocated;
+VALUE SYM_buffer_space_commited;
 
 VALUE SYM_wd;
 VALUE SYM_mask;
@@ -374,7 +380,7 @@ VALUE UM_periodically(VALUE self, VALUE interval) {
  * which the data will be read. A negative `buffer_offset` denotes a position
  * relative to the end of the buffer, e.g. a value of `-1` means the data will
  * be appended to the buffer.
- * 
+ *
  * - https://www.man7.org/linux/man-pages/man2/read.2.html
  * - https://www.man7.org/linux/man-pages/man3/io_uring_prep_read.3.html
  *
@@ -1521,17 +1527,23 @@ void Init_UM(void) {
 
   um_define_net_constants(cUM);
 
-  SYM_size =            ID2SYM(rb_intern("size"));
-  SYM_total_ops =       ID2SYM(rb_intern("total_ops"));
-  SYM_total_switches =  ID2SYM(rb_intern("total_switches"));
-  SYM_total_waits =     ID2SYM(rb_intern("total_waits"));
-  SYM_ops_pending =     ID2SYM(rb_intern("ops_pending"));
-  SYM_ops_unsubmitted = ID2SYM(rb_intern("ops_unsubmitted"));
-  SYM_ops_runqueue =    ID2SYM(rb_intern("ops_runqueue"));
-  SYM_ops_free =        ID2SYM(rb_intern("ops_free"));
-  SYM_ops_transient =   ID2SYM(rb_intern("ops_transient"));
-  SYM_time_total_cpu =  ID2SYM(rb_intern("time_total_cpu"));
-  SYM_time_total_wait = ID2SYM(rb_intern("time_total_wait"));
+  SYM_size                    = ID2SYM(rb_intern("size"));
+  SYM_total_ops               = ID2SYM(rb_intern("total_ops"));
+  SYM_total_switches          = ID2SYM(rb_intern("total_switches"));
+  SYM_total_waits             = ID2SYM(rb_intern("total_waits"));
+  SYM_ops_pending             = ID2SYM(rb_intern("ops_pending"));
+  SYM_ops_unsubmitted         = ID2SYM(rb_intern("ops_unsubmitted"));
+  SYM_ops_runqueue            = ID2SYM(rb_intern("ops_runqueue"));
+  SYM_ops_free                = ID2SYM(rb_intern("ops_free"));
+  SYM_ops_transient           = ID2SYM(rb_intern("ops_transient"));
+  SYM_time_total_cpu          = ID2SYM(rb_intern("time_total_cpu"));
+  SYM_time_total_wait         = ID2SYM(rb_intern("time_total_wait"));
+  SYM_buffer_groups           = ID2SYM(rb_intern("buffer_groups"));
+  SYM_buffers_allocated       = ID2SYM(rb_intern("buffers_allocated"));
+  SYM_buffers_free            = ID2SYM(rb_intern("buffers_free"));
+  SYM_segments_free           = ID2SYM(rb_intern("segments_free"));
+  SYM_buffer_space_allocated  = ID2SYM(rb_intern("buffer_space_allocated"));
+  SYM_buffer_space_commited   = ID2SYM(rb_intern("buffer_space_commited"));
 
   SYM_wd              = ID2SYM(rb_intern("wd"));
   SYM_mask            = ID2SYM(rb_intern("mask"));

data/ext/um/um_op.c

@@ -114,42 +114,56 @@ inline void um_op_list_compact(struct um *machine, struct um_op *head) {
   }
 }
 
+#define UM_OP_BATCH_ALLOC_SIZE (sizeof(struct um_op_result) * UM_OP_RESULT_ALLOC_BATCH_SIZE)
+
 inline struct um_op_result *multishot_result_alloc(struct um *machine) {
   if (machine->result_freelist) {
     struct um_op_result *result = machine->result_freelist;
     machine->result_freelist = result->next;
+    result->segment = NULL;
+    result->next = NULL;
     return result;
   }
 
-  struct um_op_result *batch = malloc(sizeof(struct um_op_result) * UM_OP_RESULT_ALLOC_BATCH_SIZE);
+  struct um_op_result *batch = malloc(UM_OP_BATCH_ALLOC_SIZE);
+  memset(batch, 0, UM_OP_BATCH_ALLOC_SIZE);
   for (int i = 1; i < (UM_OP_RESULT_ALLOC_BATCH_SIZE - 1); i++) {
-    batch[i].next = &batch[i + 1];
+    batch[i].next = batch + i + 1;
   }
   machine->result_freelist = batch + 1;
   return batch;
 }
 
 inline void multishot_result_free(struct um *machine, struct um_op_result *result) {
+  if (result->segment) {
+    um_segment_free(machine, result->segment);
+    result->segment = NULL;
+  }
+
   result->next = machine->result_freelist;
   machine->result_freelist = result;
 }
 
 inline void um_op_multishot_results_push(struct um *machine, struct um_op *op, __s32 res, __u32 flags) {
-  if (!op->multishot_result_count) {
-    op->result.res    = res;
-    op->result.flags  = flags;
-    op->result.next   = NULL;
-    op->multishot_result_tail = &op->result;
-  }
+  struct um_op_result *result;
+
+  if (!op->multishot_result_count)
+    result = &op->result;
   else {
-    struct um_op_result *result = multishot_result_alloc(machine);
-    result->res   = res;
-    result->flags = flags;
-    result->next  = NULL;
+    result = multishot_result_alloc(machine);
     op->multishot_result_tail->next = result;
-    op->multishot_result_tail = result;
   }
+
+  result->res   = res;
+  result->flags = flags;
+  result->next  = NULL;
+  if (op->flags & OP_F_BUFFER_POOL && res >= 0) {
+    result->segment = bp_get_op_result_segment(machine, op, res, flags);
+  }
+
+  op->multishot_result_tail = result;
   op->multishot_result_count++;
+
 }
 
 inline void um_op_multishot_results_clear(struct um *machine, struct um_op *op) {
@@ -161,6 +175,7 @@ inline void um_op_multishot_results_clear(struct um *machine, struct um_op *op)
     multishot_result_free(machine, result);
     result = next;
   }
+  op->result.next = NULL;
   op->multishot_result_tail = NULL;
   op->multishot_result_count = 0;
 }
@@ -177,6 +192,7 @@ inline struct um_op *um_op_alloc(struct um *machine) {
   for (int i = 1; i < (UM_OP_ALLOC_BATCH_SIZE - 1); i++) {
     batch[i].next = &batch[i + 1];
   }
+  batch[UM_OP_ALLOC_BATCH_SIZE - 1].next = NULL;
   machine->op_freelist = batch + 1;
   machine->metrics.ops_free += (UM_OP_ALLOC_BATCH_SIZE - 1);
   return batch;