uringmachine 0.28.3 → 0.29.1

data/ext/um/um_stream.c

@@ -1,164 +1,428 @@
 #include "um.h"
 #include <stdlib.h>
 
-static inline void stream_check_truncate_buffer(struct um_stream *stream) {
-  if ((stream->pos == stream->len) && (stream->len >= 1 << 12)) {
-    rb_str_modify(stream->buffer);
-    rb_str_set_len(stream->buffer, 0);
-    stream->len = 0;
-    stream->pos = 0;
-  }
-  else if (stream->pos >= 1 << 12) {
-    rb_str_modify(stream->buffer);
-    char *base = RSTRING_PTR(stream->buffer);
-    int len_rest = stream->len - stream->pos;
-    memmove(base, base + stream->pos, len_rest);
-    rb_str_set_len(stream->buffer, len_rest);
-    stream->len = len_rest;
-    stream->pos = 0;
+inline void stream_add_segment(struct um_stream *stream, struct um_segment *segment) {
+  segment->next = NULL;
+  if (stream->tail) {
+    stream->tail->next = segment;
+    stream->tail = segment;
   }
+  else
+    stream->head = stream->tail = segment;
+  stream->pending_len += segment->len;
 }
 
-// returns true if eof
-int stream_read_more(struct um_stream *stream) {
-  stream_check_truncate_buffer(stream);
-
-  size_t maxlen = 1 << 12;
-  size_t capa = rb_str_capacity(stream->buffer);
-  if (capa - stream->pos < maxlen)
-    rb_str_modify_expand(stream->buffer, maxlen - (capa - stream->pos));
+inline int stream_process_op_result(struct um_stream *stream, struct um_op_result *result) {
+  if (likely(result->res > 0)) {
+    if (likely(result->segment)) {
+      stream_add_segment(stream, result->segment);
+      result->segment = NULL;
+    }
+  }
+  else
+    stream->eof = 1;
 
-  rb_str_modify(stream->buffer);
-  char *ptr = RSTRING_PTR(stream->buffer) + stream->pos;
-  size_t ret = um_read_raw(stream->machine, stream->fd, ptr, maxlen);
+  return result->res;
+}
 
-  if (ret == 0) {
-    stream->eof = 1;
-    return 0;
+#define STREAM_OP_FLAGS (OP_F_MULTISHOT | OP_F_BUFFER_POOL)
+
+void stream_multishot_op_start(struct um_stream *stream) {
+  if (!stream->op)
+    stream->op = um_op_acquire(stream->machine);
+  struct io_uring_sqe *sqe;
+
+  bp_ensure_commit_level(stream->machine);
+
+  switch (stream->mode) {
+    case STREAM_BP_READ:
+      um_prep_op(stream->machine, stream->op, OP_READ_MULTISHOT, 2, STREAM_OP_FLAGS);
+      sqe = um_get_sqe(stream->machine, stream->op);
+      io_uring_prep_read_multishot(sqe, stream->fd, 0, -1, BP_BGID);
+      break;
+    case STREAM_BP_RECV:
+      um_prep_op(stream->machine, stream->op, OP_RECV_MULTISHOT, 2, STREAM_OP_FLAGS);
+      sqe = um_get_sqe(stream->machine, stream->op);
+      io_uring_prep_recv_multishot(sqe, stream->fd, NULL, 0, 0);
+	    sqe->buf_group = BP_BGID;
+	    sqe->flags |= IOSQE_BUFFER_SELECT;
+      break;
+    default:
+      um_raise_internal_error("Invalid multishot op");
   }
-
-  stream->len = stream->pos + ret;
-  rb_str_set_len(stream->buffer, stream->len);
-  return 1;
+  stream->op->bp_commit_level = stream->machine->bp_commit_level;
 }
 
-// ensures given string can hold at least given len bytes (+trailing null)
-static inline void str_expand(VALUE str, size_t len) {
-  rb_str_resize(str, len);
+void stream_multishot_op_stop(struct um_stream *stream) {
+  assert(!stream->op);
+
+  if (!(stream->op->flags & OP_F_CQE_DONE)) {
+    stream->op->flags |= OP_F_ASYNC;
+    um_cancel_op(stream->machine, stream->op);
+  }
+  else
+    um_op_release(stream->machine, stream->op);
+  stream->op = NULL;
 }
 
-static inline void str_copy_bytes(VALUE dest, const char *src, ssize_t len) {
-  str_expand(dest, len + 1);
-  char *dest_ptr = RSTRING_PTR(dest);
-  memcpy(dest_ptr, src, len);
-  dest_ptr[len] = 0;
-  rb_str_set_len(dest, len);
+void um_stream_cleanup(struct um_stream *stream) {
+  if (stream->op) stream_multishot_op_stop(stream);
+
+  while (stream->head) {
+    struct um_segment *next = stream->head->next;
+    um_segment_checkin(stream->machine, stream->head);
+    stream->head = next;
+  }
+  stream->pending_len = 0;
 }
 
-VALUE stream_get_line(struct um_stream *stream, VALUE buf, ssize_t maxlen) {
-  while (true) {
-    char *start = RSTRING_PTR(stream->buffer) + stream->pos;
-    ssize_t pending_len = stream->len - stream->pos;
-    ssize_t search_len = pending_len;
-    ssize_t absmax_len = labs(maxlen);
-    int should_limit_len = (absmax_len > 0) && (search_len > maxlen);
-    if (should_limit_len) search_len = absmax_len;
-
-    char * lf_ptr = memchr(start, '\n', search_len);
-    if (lf_ptr) {
-      ssize_t len = lf_ptr - start;
-      if (len && (start[len - 1] == '\r')) len -= 1;
+// returns true if case of ENOBUFS error, sets more to true if more data forthcoming
+inline int stream_process_segments(
+  struct um_stream *stream, size_t *total_bytes, int *more) {
+
+  *more = 0;
+  struct um_op_result *result = &stream->op->result;
+  stream->op->flags &= ~OP_F_CQE_SEEN;
+  while (result) {
+    if (unlikely(result->res == -ENOBUFS)) {
+      *more = 0;
+      return true;
+    }
+    if (unlikely(result->res == -ECANCELED)) {
+      *more = 0;
+      return false;
+    }
+    um_raise_on_error_result(result->res);
 
-      stream->pos += lf_ptr - start + 1;
-      if (NIL_P(buf)) return rb_utf8_str_new(start, len);
+    *more = (result->flags & IORING_CQE_F_MORE);
+    *total_bytes += result->res;
+    stream_process_op_result(stream, result);
+    result = result->next;
+  }
+  return false;
+}
 
-      str_copy_bytes(buf, start, len);
-      return buf;
+void stream_clear(struct um_stream *stream) {
+  if (stream->op && stream->machine->ring_initialized) {
+    if (OP_CQE_SEEN_P(stream->op)) {
+      size_t total_bytes = 0;
+      int more = false;
+      stream_process_segments(stream, &total_bytes, &more);
+      um_op_multishot_results_clear(stream->machine, stream->op);
     }
-    else if (should_limit_len && pending_len > search_len)
-      // hit maxlen
-      return Qnil;
 
-    if (!stream_read_more(stream))
-      return Qnil;
+    if (OP_CQE_DONE_P(stream->op))
+      um_op_release(stream->machine, stream->op);
+    else
+      um_cancel_op_and_discard_cqe(stream->machine, stream->op);
+
+    stream->op = NULL;
+  }
+
+  while (stream->head) {
+    struct um_segment *next = stream->head->next;
+    um_segment_checkin(stream->machine, stream->head);
+    stream->head = next;
+  }
+  stream->pending_len = 0;
+
+  if (stream->working_buffer) {
+    bp_buffer_checkin(stream->machine, stream->working_buffer);
+    stream->working_buffer = NULL;
   }
 }
 
-VALUE stream_get_string(struct um_stream *stream, VALUE buf, ssize_t len) {
-  size_t abslen = labs(len);
-  while (stream->len - stream->pos < abslen) {
-    if (!stream_read_more(stream)) {
-      if (len > 0)
-        return Qnil;
+inline void stream_await_segments(struct um_stream *stream) {
+  if (unlikely(!stream->op)) stream_multishot_op_start(stream);
+
+  if (!OP_CQE_SEEN_P(stream->op)) {
+    stream->op->flags &= ~OP_F_ASYNC;
+    VALUE ret = um_yield(stream->machine);
+    stream->op->flags |= OP_F_ASYNC;
+    if (!OP_CQE_SEEN_P(stream->op)) RAISE_IF_EXCEPTION(ret);
+    RB_GC_GUARD(ret);
+  }
+}
+
+int stream_get_more_segments_bp(struct um_stream *stream) {
+  size_t total_bytes = 0;
+  int more = false;
+  int enobufs = false;
+
+  while (1) {
+    if (unlikely(stream->eof)) return 0;
+
+    stream_await_segments(stream);
+    enobufs = stream_process_segments(stream, &total_bytes, &more);
+    um_op_multishot_results_clear(stream->machine, stream->op);
+    if (unlikely(enobufs)) {
+      int should_restart = stream->pending_len < (stream->machine->bp_buffer_size * 4);
+      // int same_threshold = stream->op->bp_commit_level == stream->machine->bp_commit_level;
+
+      // fprintf(stderr, "%p enobufs total: %ld pending: %ld threshold: %ld bc: %d (same: %d, restart: %d)\n",
+      //   stream,
+      //   total_bytes, stream->pending_len, stream->machine->bp_commit_level,
+      //   stream->machine->bp_buffer_count,
+      //   same_threshold, should_restart
+      // );
+
+      // If multiple stream ops are happening at the same time, they'll all get
+      // ENOBUFS! We track the commit threshold in the op in order to prevent
+      // running bp_handle_enobufs() more than once.
+
+      if (should_restart) {
+        if (stream->op->bp_commit_level == stream->machine->bp_commit_level)
+          bp_handle_enobufs(stream->machine);
+  
+        um_op_release(stream->machine, stream->op);
+        stream->op = NULL;
+        // stream_multishot_op_start(stream);
+      }
+      else {
+        um_op_release(stream->machine, stream->op);
+        stream->op = NULL;
+      }
 
-      abslen = stream->len - stream->pos;
+      if (total_bytes) return total_bytes;
+    }
+    else {
+      if (more)
+        stream->op->flags &= ~OP_F_CQE_SEEN;
+      if (total_bytes || stream->eof) return total_bytes;
     }
   }
+}
 
-  char *start = RSTRING_PTR(stream->buffer) + stream->pos;
-  stream->pos += abslen;
+int stream_get_more_segments_ssl(struct um_stream *stream) {
+  if (!stream->working_buffer)
+    stream->working_buffer = bp_buffer_checkout(stream->machine);
 
-  if (NIL_P(buf)) return rb_utf8_str_new(start, abslen);
+  char *ptr = stream->working_buffer->buf + stream->working_buffer->pos;
+  size_t maxlen = stream->working_buffer->len - stream->working_buffer->pos;
+  int res = um_ssl_read_raw(stream->machine, stream->target, ptr, maxlen);
+  if (res == 0) return 0;
+  if (res < 0) rb_raise(eUMError, "Failed to read segment");
 
-  str_copy_bytes(buf, start, abslen);
-  return buf;
+  struct um_segment *segment = bp_buffer_consume(stream->machine, stream->working_buffer, res);
+  if ((size_t)res == maxlen) {
+    bp_buffer_checkin(stream->machine, stream->working_buffer);
+    stream->working_buffer = NULL;
+  }
+  stream_add_segment(stream, segment);
+  return 1;
 }
 
-VALUE stream_skip(struct um_stream *stream, size_t len) {
-  while (stream->len - stream->pos < len)
-    if (!stream_read_more(stream)) {
-      return Qnil;
+int stream_get_more_segments(struct um_stream *stream) {
+  switch (stream->mode) {
+    case STREAM_BP_READ:
+    case STREAM_BP_RECV:
+      return stream_get_more_segments_bp(stream);
+    case STREAM_SSL:
+      return stream_get_more_segments_ssl(stream);
+    default:
+      rb_raise(eUMError, "Invalid stream mode");
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////////////
+
+inline void stream_shift_head(struct um_stream *stream) {
+  struct um_segment *consumed = stream->head;
+  stream->head = consumed->next;
+  if (!stream->head) stream->tail = NULL;
+  um_segment_checkin(stream->machine, consumed);
+  stream->pos = 0;
+}
+
+inline void stream_skip(struct um_stream *stream, size_t inc, int safe_inc) {
+  while (inc) {
+    size_t segment_len = stream->head->len - stream->pos;
+    size_t inc_len = (segment_len <= inc) ? segment_len : inc;
+    inc -= inc_len;
+    stream->pos += inc_len;
+    stream->pending_len -= inc_len;
+    if (stream->pos == stream->head->len) {
+      stream_shift_head(stream);
+      if (inc && safe_inc && !stream->head) {
+        if (!stream_get_more_segments(stream)) break;
+      }
     }
+  }
+}
 
-  stream->pos += len;
-  return NUM2INT(len);
+inline void stream_copy(struct um_stream *stream, char *dest, size_t len) {
+  while (len) {
+    char *segment_ptr = stream->head->ptr + stream->pos;
+    size_t segment_len = stream->head->len - stream->pos;
+    size_t cpy_len = (segment_len <= len) ? segment_len : len;
+    memcpy(dest, segment_ptr, cpy_len);
+
+    len -= cpy_len;
+    stream->pos += cpy_len;
+    stream->pending_len -= cpy_len;
+    dest += cpy_len;
+    if (stream->pos == stream->head->len) stream_shift_head(stream);
+  }
 }
 
-VALUE resp_get_line(struct um_stream *stream, VALUE out_buffer) {
-  char *start = RSTRING_PTR(stream->buffer) + stream->pos;
+VALUE stream_consume_string(struct um_stream *stream, VALUE out_buffer, size_t len, size_t inc, int safe_inc) {
+  VALUE str = Qnil;
+  if (!NIL_P(out_buffer)) {
+    str = out_buffer;
+    size_t str_len = RSTRING_LEN(str);
+    if (str_len < len)
+      rb_str_resize(str, len);
+    else if (str_len > len)
+      rb_str_set_len(str, len);
+  }
+  else
+    str = rb_str_new(NULL, len);
+  char *dest = RSTRING_PTR(str);
+
+  stream_copy(stream, dest, len);
+  stream_skip(stream, inc, safe_inc);
+  return str;
+  RB_GC_GUARD(str);
+}
+
+VALUE stream_get_line(struct um_stream *stream, VALUE out_buffer, size_t maxlen) {
+  // if (stream->head) {
+  //   fprintf(stderr, "head: %p pos: %ld: %.*s\n",
+  //     stream->head, stream->pos,
+  //     (int)(stream->head->len - stream->pos), stream->head->ptr + stream->pos
+  //   );
+  // }
+
+  if (unlikely(stream->eof && !stream->head)) return Qnil;
+  if (!stream->tail && !stream_get_more_segments(stream)) return Qnil;
+
+  struct um_segment *last = NULL;
+  struct um_segment *current = stream->head;
+  size_t remaining_len = maxlen;
+  size_t total_len = 0;
+  size_t inc = 1;
+  size_t pos = stream->pos;
+
   while (true) {
-    char * lf_ptr = memchr(start, '\r', stream->len - stream->pos);
+    size_t segment_len = current->len - pos;
+    size_t search_len = segment_len;
+    if (maxlen && (search_len > remaining_len)) search_len = remaining_len;
+    char *start = current->ptr + pos;
+    char *lf_ptr = memchr(start, '\n', search_len);
+
     if (lf_ptr) {
-      ulong len = lf_ptr - start;
-      stream->pos += len + 2;
-
-      if (NIL_P(out_buffer)) {
-        VALUE str = rb_interned_str(start, len + 1);
-        rb_str_set_len(str, len);
-        RSTRING_PTR(str)[len] = 0;
-        RB_GC_GUARD(str);
-        return str;
+      size_t len = lf_ptr - start;
+
+      total_len += len;
+
+      // search for \r
+      if (total_len) {
+        if (len) {
+          if  (start[len - 1] == '\r') {
+            total_len -= 1;
+            inc = 2;
+          }
+        }
+        else {
+          if (last && (((char *)last->ptr)[last->len - 1] == '\r')) {
+            total_len -= 1;
+            inc = 2;
+          }
+        }
       }
 
-      str_copy_bytes(out_buffer, start, len);
-      return out_buffer;
+      return stream_consume_string(stream, out_buffer, total_len, inc, false);
     }
+    else {
+      // not found, early return if segment len exceeds maxlen
+      if (maxlen && segment_len >= maxlen) return Qnil;
 
-    if (stream_read_more(stream))
-      // buffer ptr and pos may have changed after reading
-      start = RSTRING_PTR(stream->buffer) + stream->pos;
-    else
-      return Qnil;
+      total_len += segment_len;
+      remaining_len -= segment_len;
+    }
+
+    if (!current->next) {
+      if (!stream_get_more_segments(stream)) {
+        return Qnil;
+      }
+    }
+
+    last = current;
+    current = current->next;
+    pos = 0;
   }
 }
 
-VALUE resp_get_string(struct um_stream *stream, ulong len, VALUE out_buffer) {
-  ulong read_len = len + 2;
+VALUE stream_get_string(struct um_stream *stream, VALUE out_buffer, ssize_t len, size_t inc, int safe_inc) {
+  if (unlikely(stream->eof && !stream->head)) return Qnil;
+  if (!stream->tail && !stream_get_more_segments(stream)) return Qnil;
 
-  while (stream->len - stream->pos < read_len)
-    if (!stream_read_more(stream)) return Qnil;
+  struct um_segment *current = stream->head;
+  size_t abs_len = labs(len);
+  size_t remaining_len = abs_len;
+  size_t total_len = 0;
+  size_t pos = stream->pos;
 
-  char *start = RSTRING_PTR(stream->buffer) + stream->pos;
-  stream->pos += read_len;
+  while (true) {
+    size_t segment_len = current->len - pos;
+    if (abs_len && segment_len > remaining_len) {
+      segment_len = remaining_len;
+    }
+    total_len += segment_len;
+    if (abs_len) {
+      remaining_len -= segment_len;
+      if (!remaining_len)
+        return stream_consume_string(stream, out_buffer, total_len, inc, safe_inc);
+    }
+
+    if (!current->next) {
+      if (len <= 0)
+        return stream_consume_string(stream, out_buffer, total_len, inc, safe_inc);
+
+      if (!stream_get_more_segments(stream))
+        return Qnil;
+    }
+    current = current->next;
+    pos = 0;
+  }
+}
+
+VALUE resp_get_line(struct um_stream *stream, VALUE out_buffer) {
+  if (unlikely(stream->eof && !stream->head)) return Qnil;
+  if (!stream->tail && !stream_get_more_segments(stream)) return Qnil;
+
+  struct um_segment *current = stream->head;
+  size_t total_len = 0;
+  size_t pos = stream->pos;
+
+  while (true) {
+    size_t segment_len = current->len - pos;
+    char *start = current->ptr + pos;
+    char *lf_ptr = memchr(start, '\r', segment_len);
+    if (lf_ptr) {
+      size_t len = lf_ptr - start;
+      total_len += len;
+      return stream_consume_string(stream, out_buffer, total_len, 2, true);
+    }
+    else
+      total_len += segment_len;
 
-  if (NIL_P(out_buffer)) return rb_utf8_str_new(start, len);
+    if (!current->next)
+      if (!stream_get_more_segments(stream)) return Qnil;
 
-  str_copy_bytes(out_buffer, start, len);
-  return out_buffer;
+    current = current->next;
+  }
+}
+
+inline VALUE resp_get_string(struct um_stream *stream, ulong len, VALUE out_buffer) {
+  return stream_get_string(stream, out_buffer, len, 2, true);
 }
 
 inline ulong resp_parse_length_field(const char *ptr, int len) {
-  return strtoul(ptr + 1, NULL, 10);
+  ulong acc = 0;
+  for(int i = 1; i < len; i++)
+    acc = acc * 10 + (ptr[i] - '0');
+  return acc;
 }
 
 VALUE resp_decode_hash(struct um_stream *stream, VALUE out_buffer, ulong len) {
@@ -180,8 +444,7 @@ VALUE resp_decode_array(struct um_stream *stream, VALUE out_buffer, ulong len) {
   VALUE array = rb_ary_new2(len);
 
   for (ulong i = 0; i < len; i++) {
-    VALUE buf = rb_str_new(NULL, 100);
-    VALUE value = resp_decode(stream, buf);
+    VALUE value = resp_decode(stream, out_buffer);
     rb_ary_push(array, value);
     RB_GC_GUARD(value);
   }
@@ -191,11 +454,11 @@ VALUE resp_decode_array(struct um_stream *stream, VALUE out_buffer, ulong len) {
 }
 
 static inline VALUE resp_decode_simple_string(char *ptr, ulong len) {
-  return rb_interned_str(ptr + 1, len - 1);
+  return rb_str_new(ptr + 1, len - 1);
 }
 
-static inline VALUE resp_decode_string(struct um_stream *stream, ulong len) {
-  return resp_get_string(stream, len, Qnil);
+static inline VALUE resp_decode_string(struct um_stream *stream, VALUE out_buffer, ulong len) {
+  return resp_get_string(stream, len, out_buffer);
 }
 
 static inline VALUE resp_decode_string_with_encoding(struct um_stream *stream, VALUE out_buffer, ulong len) {
@@ -221,7 +484,7 @@ static inline VALUE resp_decode_simple_error(char *ptr, ulong len) {
   static ID ID_new = 0;
   if (!ID_new) ID_new = rb_intern("new");
 
-  VALUE msg = rb_interned_str(ptr + 1, len - 1);
+  VALUE msg = rb_str_new(ptr + 1, len - 1);
   VALUE err = rb_funcall(eStreamRESPError, ID_new, 1, msg);
   RB_GC_GUARD(msg);
   return err;
@@ -231,7 +494,7 @@ static inline VALUE resp_decode_error(struct um_stream *stream, VALUE out_buffer
   static ID ID_new = 0;
   if (!ID_new) ID_new = rb_intern("new");
 
-  VALUE msg = resp_decode_string(stream, len);
+  VALUE msg = resp_decode_string(stream, out_buffer, len);
   VALUE err = rb_funcall(eStreamRESPError, ID_new, 1, msg);
   RB_GC_GUARD(msg);
   return err;
@@ -262,7 +525,7 @@ VALUE resp_decode(struct um_stream *stream, VALUE out_buffer) {
       return resp_decode_simple_string(ptr, len);
     case '$': // string
       data_len = resp_parse_length_field(ptr, len);
-      return resp_decode_string(stream, data_len);
+      return resp_decode_string(stream, out_buffer, data_len);
     case '=': // string with encoding
       data_len = resp_parse_length_field(ptr, len);
       return resp_decode_string_with_encoding(stream, out_buffer, data_len);
@@ -404,34 +667,3 @@ void resp_encode(struct um_write_buffer *buf, VALUE obj) {
       um_raise_internal_error("Can't encode object");
   }
 }
-
-void resp_encode_cmd(struct um_write_buffer *buf, int argc, VALUE *argv) {
-  char tmp1[48];
-  char tmp2[60];
-
-  sprintf(tmp1, "*%d\r\n", argc);
-  write_buffer_append_cstr(buf, tmp1);
-  for (int i = 0; i < argc; i++) {
-    switch (TYPE(argv[i])) {
-      case T_FIXNUM:
-        sprintf(tmp1, "%ld", NUM2LONG(argv[i]));
-        sprintf(tmp2, "$%ld\r\n%s\r\n", strlen(tmp1), (char *)tmp1);
-        write_buffer_append_cstr(buf, tmp2);
-        break;
-      case T_FLOAT:
-        sprintf(tmp1, "%lg", NUM2DBL(argv[i]));
-        sprintf(tmp2, "$%ld\r\n%s\r\n", strlen(tmp1), (char *)tmp1);
-        write_buffer_append_cstr(buf, tmp2);
-        break;
-      case T_STRING:
-        write_buffer_append_resp_bulk_string(buf, argv[i]);
-        break;
-      case T_SYMBOL:
-        write_buffer_append_resp_bulk_string(buf, rb_sym_to_s(argv[i]));
-        break;
-      default:
-        um_raise_internal_error("Can't encode object");
-    }
-  }
-  return;
-}