uringmachine 0.29.2 → 0.31.0

data/ext/um/um_connection_class.c

@@ -0,0 +1,394 @@
+#include "um.h"
+
+VALUE cConnection;
+VALUE eConnectionRESPError;
+
+VALUE SYM_fd;
+VALUE SYM_socket;
+VALUE SYM_ssl;
+
+inline int connection_has_target_obj_p(struct um_connection *conn) {
+  switch (conn->mode) {
+    case CONNECTION_SSL:
+    case CONNECTION_STRING:
+    case CONNECTION_IO_BUFFER:
+      return true;
+    default:
+      return false;
+  }
+}
+
+inline void connection_mark_segments(struct um_connection *conn) {
+  struct um_segment *curr = conn->head;
+  while (curr) {
+    // rb_gc_mark_movable(curr->obj);
+    curr = curr->next;
+  }
+}
+
+inline void connection_compact_segments(struct um_connection *conn) {
+  struct um_segment *curr = conn->head;
+  while (curr) {
+    // curr->obj = rb_gc_location(curr->obj);
+    curr = curr->next;
+  }
+}
+
+static void Connection_mark(void *ptr) {
+  struct um_connection *conn = ptr;
+  rb_gc_mark_movable(conn->self);
+  rb_gc_mark_movable(conn->machine->self);
+
+  if (connection_has_target_obj_p(conn)) {
+    rb_gc_mark_movable(conn->target);
+    connection_mark_segments(conn);
+  }
+}
+
+static void Connection_compact(void *ptr) {
+  struct um_connection *conn = ptr;
+  conn->self = rb_gc_location(conn->self);
+
+  if (connection_has_target_obj_p(conn)) {
+    conn->target = rb_gc_location(conn->target);
+    connection_compact_segments(conn);
+  }
+}
+
+static void Connection_free(void *ptr) {
+  struct um_connection *conn = ptr;
+  connection_clear(conn);
+}
+
+static const rb_data_type_t Connection_type = {
+  .wrap_struct_name = "UringMachine::Connection",
+  .function = {
+    .dmark = Connection_mark,
+    .dfree = Connection_free,
+    .dsize = NULL,
+    .dcompact = Connection_compact
+  },
+  .flags = RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_WB_PROTECTED | RUBY_TYPED_EMBEDDABLE
+};
+
+static VALUE Connection_allocate(VALUE klass) {
+  struct um_connection *conn;
+  VALUE self = TypedData_Make_Struct(klass, struct um_connection, &Connection_type, conn);
+  return self;
+}
+
+static inline struct um_connection *um_get_connection(VALUE self) {
+  struct um_connection *conn;
+  TypedData_Get_Struct(self, struct um_connection, &Connection_type, conn);
+  return conn;
+}
+
+static inline void connection_set_target(struct um_connection *conn, VALUE target, enum um_connection_mode mode) {
+  conn->mode = mode;
+  switch (mode) {
+    case CONNECTION_FD:
+    case CONNECTION_SOCKET:
+      conn->fd = NUM2INT(target);
+      return;
+    case CONNECTION_SSL:
+      conn->target = target;
+      um_ssl_set_bio(conn->machine, target);
+      return;
+    default:
+      rb_raise(eUMError, "Invalid connection mode");
+  }
+}
+
+static inline void connection_setup(struct um_connection *conn, VALUE target, VALUE mode) {
+  conn->working_buffer = NULL;
+  if (NIL_P(mode)) {
+    if (TYPE(target) == T_DATA)
+      connection_set_target(conn, target, CONNECTION_SSL);
+    else
+      connection_set_target(conn, target, CONNECTION_FD);
+  }
+  else if (mode == SYM_fd)
+    connection_set_target(conn, target, CONNECTION_FD);
+  else if (mode == SYM_socket)
+    connection_set_target(conn, target, CONNECTION_SOCKET);
+  else if (mode == SYM_ssl)
+    connection_set_target(conn, target, CONNECTION_SSL);
+  else
+    rb_raise(eUMError, "Invalid connection mode");
+}
+
+/* call-seq:
+ *   UM::Stream.new(machine, fd, mode = nil) -> conn
+ *   machine.connection(fd, mode = nil) -> conn
+ *   machine.connection(fd, mode = nil) { |conn| ... }
+ *
+ * Initializes a new connection with the given UringMachine instance, target and
+ * optional mode. The target maybe a file descriptor, or an instance of
+ * OpenSSL::SSL::SSLSocket. In case of an SSL socket, the mode should be :ssl.
+ *
+ * @param machine [UringMachine] UringMachine instance
+ * @param target [integer, OpenSSL::SSL::SSLSocket] connection target: file descriptor or SSL socket
+ * @param mode [Symbol] optional connection mode: :fd, :socket, :ssl
+ * @return [void]
+ */
+VALUE Connection_initialize(int argc, VALUE *argv, VALUE self) {
+  VALUE machine;
+  VALUE target;
+  VALUE mode;
+  rb_scan_args(argc, argv, "21", &machine, &target, &mode);
+
+  struct um_connection *conn = um_get_connection(self);
+  memset(conn, 0, sizeof(struct um_connection));
+
+  RB_OBJ_WRITE(self, &conn->self, self);
+  conn->machine = um_get_machine(machine);
+  connection_setup(conn, target, mode);
+
+  return self;
+}
+
+/* call-seq:
+ *   conn.mode -> mode
+ *
+ * Returns the connection mode.
+ *
+ * @return [Symbol] connection mode
+ */
+VALUE Connection_mode(VALUE self) {
+  struct um_connection *conn = um_get_connection(self);
+  switch (conn->mode) {
+    case CONNECTION_FD:  return SYM_fd;
+    case CONNECTION_SOCKET:  return SYM_socket;
+    case CONNECTION_SSL:      return SYM_ssl;
+    default:              return Qnil;
+  }
+  return Qnil;
+}
+
+/* call-seq:
+ *   conn.read_line(limit) -> str
+ *
+ * Reads from the string until a newline character is encountered. Returns the
+ * line without the newline delimiter. If limit is 0, the line length is not
+ * limited. If no newline delimiter is found before EOF, returns nil.
+ *
+ * @param limit [integer] maximum line length (0 means no limit)
+ * @return [String, nil] read data or nil
+ */
+VALUE Connection_read_line(VALUE self, VALUE limit) {
+  struct um_connection *conn = um_get_connection(self);
+  return connection_read_line(conn, Qnil, NUM2ULONG(limit));
+}
+
+/* call-seq:
+ *   conn.read(len) -> str
+ *
+ * Reads len bytes from the conn. If len is 0, reads all available bytes. If
+ * len is negative, reads up to -len available bytes. If len is positive and eof
+ * is encountered before len bytes are read, returns nil.
+ *
+ * @param len [integer] number of bytes to read
+ * @return [String, nil] read data or nil
+ */
+VALUE Connection_read(VALUE self, VALUE len) {
+  struct um_connection *conn = um_get_connection(self);
+  return connection_read(conn, Qnil, NUM2LONG(len), 0, false);
+}
+
+/* call-seq:
+ *   conn.read_to_delim(delim, limit) -> str
+ *
+ * Reads from the string until a the given delimiter is encountered. Returns the
+ * line without the delimiter. If limit is 0, the length is not limited. If a
+ * delimiter is not found before EOF and limit is 0 or greater, returns nil.
+ *
+ * If no delimiter is found before EOF and limit is negative, returns the
+ * buffered data up to EOF or until the absolute-value length limit is reached.
+ *
+ * The `delim` parameter must be a single byte string.
+ *
+ * @param delim [String] delimiter (single byte) @param limit [integer] maximum
+ * line length (0 means no limit) @return [String, nil] read data or nil
+ */
+VALUE Connection_read_to_delim(VALUE self, VALUE delim, VALUE limit) {
+  struct um_connection *conn = um_get_connection(self);
+  return connection_read_to_delim(conn, Qnil, delim, NUM2LONG(limit));
+}
+
+/* call-seq:
+ *   conn.skip(len) -> len
+ *
+ * Skips len bytes in the conn.
+ *
+ * @param len [integer] number of bytes to skip
+ * @return [Integer] len
+ */
+VALUE Connection_skip(VALUE self, VALUE len) {
+  struct um_connection *conn = um_get_connection(self);
+  connection_skip(conn, NUM2LONG(len), true);
+  return len;
+}
+
+/* call-seq:
+ *   conn.read_each { |data| } -> conn
+ *
+ * Reads from the target, passing each chunk to the given block.
+ *
+ * @return [UringMachine::Connection] conn
+ */
+VALUE Connection_read_each(VALUE self) {
+  struct um_connection *conn = um_get_connection(self);
+  connection_read_each(conn);
+  return self;
+}
+
+/* call-seq:
+ *   conn.write(*bufs) -> len
+ *
+ * Writes to the connection, ensuring that all data has been written before
+ * returning the total number of bytes written.
+ *
+ * @param bufs [Array<String, IO::Buffer>] data to write
+ * @return [Integer] total bytes written
+ */
+VALUE Connection_write(int argc, VALUE *argv, VALUE self) {
+  struct um_connection *conn = um_get_connection(self);
+  return connection_writev(conn, argc, argv);
+}
+
+/* call-seq:
+ *   conn.resp_read -> obj
+ *
+ * Decodes an object from a RESP (Redis protocol) message.
+ *
+ * @return [any] decoded object
+ */
+VALUE Connection_resp_read(VALUE self) {
+  struct um_connection *conn = um_get_connection(self);
+  VALUE out_buffer = rb_utf8_str_new_literal("");
+  VALUE obj = resp_read(conn, out_buffer);
+  RB_GC_GUARD(out_buffer);
+  return obj;
+}
+
+/* call-seq:
+ *   conn.resp_write(obj) -> conn
+ *
+ * Writes the given object using RESP (Redis protocol) to the connection target.
+ * Returns the number of bytes written.
+ *
+ * @param obj [any] object to write
+ * @return [Integer] total bytes written
+ */
+VALUE Connection_resp_write(VALUE self, VALUE obj) {
+  struct um_connection *conn = um_get_connection(self);
+
+  VALUE str = rb_str_new(NULL, 0);
+  struct um_write_buffer buf;
+  write_buffer_init(&buf, str);
+  rb_str_modify(str);
+  resp_encode(&buf, obj);
+  write_buffer_update_len(&buf);
+
+  size_t len = connection_write_raw(conn, buf.ptr, buf.len);
+  RB_GC_GUARD(str);
+  return ULONG2NUM(len);
+}
+
+/* call-seq:
+ *   conn.resp_encode(obj) -> string
+ *
+ * Encodes an object into a RESP (Redis protocol) message.
+ *
+ * @param str [String] string buffer
+ * @param obj [any] object to be encoded
+ * @return [String] str
+ */
+VALUE Connection_resp_encode(VALUE self, VALUE str, VALUE obj) {
+  struct um_write_buffer buf;
+  write_buffer_init(&buf, str);
+  rb_str_modify(str);
+  resp_encode(&buf, obj);
+  write_buffer_update_len(&buf);
+  return str;
+}
+
+/* call-seq:
+ *   conn.eof? -> bool
+ *
+ * Returns true if connection has reached EOF.
+ *
+ * @return [bool] EOF reached
+ */
+VALUE Connection_eof_p(VALUE self) {
+  struct um_connection *conn = um_get_connection(self);
+  return conn->eof ? Qtrue : Qfalse;
+}
+
+/* call-seq:
+ *   conn.consumed -> int
+ *
+ * Returns the total number of bytes consumed from the conn.
+ *
+ * @return [Integer] total bytes consumed
+ */
+VALUE Connection_consumed(VALUE self) {
+  struct um_connection *conn = um_get_connection(self);
+  return LONG2NUM(conn->consumed_bytes);
+}
+
+/* call-seq:
+ *   conn.pending -> int
+ *
+ * Returns the number of bytes available for reading.
+ *
+ * @return [Integer] bytes available
+ */
+VALUE Connection_pending(VALUE self) {
+  struct um_connection *conn = um_get_connection(self);
+  return LONG2NUM(conn->pending_bytes);
+}
+
+/* call-seq:
+ *   conn.clear -> conn
+ *
+ * Clears all available bytes and stops any ongoing read operation.
+ *
+ * @return [UM::Stream] self
+ */
+VALUE Connection_clear(VALUE self) {
+  struct um_connection *conn = um_get_connection(self);
+  connection_clear(conn);
+  return self;
+}
+
+void Init_Stream(void) {
+  cConnection = rb_define_class_under(cUM, "Connection", rb_cObject);
+  rb_define_alloc_func(cConnection, Connection_allocate);
+
+  rb_define_method(cConnection, "initialize", Connection_initialize, -1);
+  rb_define_method(cConnection, "mode", Connection_mode, 0);
+
+  rb_define_method(cConnection, "read_line", Connection_read_line, 1);
+  rb_define_method(cConnection, "read", Connection_read, 1);
+  rb_define_method(cConnection, "read_to_delim", Connection_read_to_delim, 2);
+  rb_define_method(cConnection, "skip", Connection_skip, 1);
+  rb_define_method(cConnection, "read_each", Connection_read_each, 0);
+
+  rb_define_method(cConnection, "write", Connection_write, -1);
+
+  rb_define_method(cConnection, "resp_read", Connection_resp_read, 0);
+  rb_define_method(cConnection, "resp_write", Connection_resp_write, 1);
+  rb_define_singleton_method(cConnection, "resp_encode", Connection_resp_encode, 2);
+
+  rb_define_method(cConnection, "eof?", Connection_eof_p, 0);
+  rb_define_method(cConnection, "consumed", Connection_consumed, 0);
+  rb_define_method(cConnection, "pending", Connection_pending, 0);
+  rb_define_method(cConnection, "clear", Connection_clear, 0);
+
+  eConnectionRESPError = rb_define_class_under(cConnection, "RESPError", rb_eStandardError);
+
+  SYM_fd = ID2SYM(rb_intern("fd"));
+  SYM_socket = ID2SYM(rb_intern("socket"));
+  SYM_ssl     = ID2SYM(rb_intern("ssl"));
+}

data/ext/um/um_ssl.c

@@ -8,14 +8,14 @@
 static int um_bio_read(BIO *bio, char *buf, int blen)
 {
   struct um *machine = (struct um *)BIO_get_ex_data(bio, IDX_BIO_DATA_MACHINE);
-  long fd = (long)BIO_get_ex_data(bio, IDX_BIO_DATA_FD);
+  int fd = (int)(long)BIO_get_ex_data(bio, IDX_BIO_DATA_FD);
   return (int)um_read_raw(machine, fd, buf, blen);
 }
 
 static int um_bio_write(BIO *bio, const char *buf, int blen)
 {
   struct um *machine = (struct um *)BIO_get_ex_data(bio, IDX_BIO_DATA_MACHINE);
-  long fd = (long)BIO_get_ex_data(bio, IDX_BIO_DATA_FD);
+  int fd = (int)(long)BIO_get_ex_data(bio, IDX_BIO_DATA_FD);
   return (int)um_write_raw(machine, fd, buf, blen);
 }
 
@@ -75,7 +75,7 @@ void um_ssl_set_bio(struct um *machine, VALUE ssl_obj)
   SSL_set0_wbio(ssl, bio);
 }
 
-int um_ssl_read_raw(struct um *machine, VALUE ssl_obj, char *ptr, int maxlen) {
+int um_ssl_read_raw(struct um *machine, VALUE ssl_obj, char *ptr, size_t maxlen) {
   SSL *ssl = RTYPEDDATA_GET_DATA(ssl_obj);
 
   int ret = SSL_read(ssl, ptr, maxlen);
@@ -84,23 +84,59 @@ int um_ssl_read_raw(struct um *machine, VALUE ssl_obj, char *ptr, int maxlen) {
   return ret;
 }
 
-int um_ssl_read(struct um *machine, VALUE ssl_obj, VALUE buf, int maxlen) {
+int um_ssl_read(struct um *machine, VALUE ssl_obj, VALUE buf, size_t maxlen) {
   void *ptr = um_prepare_read_buffer(buf, maxlen, 0);
   int ret = um_ssl_read_raw(machine, ssl_obj, ptr, maxlen);
   um_update_read_buffer(buf, 0, ret);
   return ret;
 }
 
-int um_ssl_write(struct um *machine, VALUE ssl_obj, VALUE buf, int len) {
+int um_ssl_write(struct um *machine, VALUE ssl_obj, VALUE buf, size_t len) {
   SSL *ssl = RTYPEDDATA_GET_DATA(ssl_obj);
   const void *base;
   size_t size;
   um_get_buffer_bytes_for_writing(buf, &base, &size, true);
-  if ((len == (int)-1) || (len > (int)size)) len = (int)size;
+  if (!len || (len > size)) len = size;
+  if (len > INT_MAX) len = INT_MAX;
   if (unlikely(!len)) return INT2NUM(0);
 
-  int ret = SSL_write(ssl, base, len);
+  int ret = SSL_write(ssl, base, (int)len);
   if (ret <= 0) rb_raise(eUMError, "Failed to write");
 
   return ret;
 }
+
+int um_ssl_write_raw(struct um *machine, VALUE ssl_obj, const char *buffer, size_t len) {
+  SSL *ssl = RTYPEDDATA_GET_DATA(ssl_obj);
+  if (unlikely(!len)) return INT2NUM(0);
+
+  int ret = SSL_write(ssl, buffer, (int)len);
+  if (ret <= 0) rb_raise(eUMError, "Failed to write");
+
+  return ret;
+}
+
+int um_ssl_write_all(struct um *machine, VALUE ssl_obj, VALUE buf) {
+  SSL *ssl = RTYPEDDATA_GET_DATA(ssl_obj);
+  const char *base;
+  size_t size;
+  um_get_buffer_bytes_for_writing(buf, (const void **)&base, &size, true);
+
+  size_t left = size;
+  while (left) {
+    int ret = SSL_write(ssl, base, (int)left);
+    if (ret <= 0) rb_raise(eUMError, "Failed to write");
+
+    left -= ret;
+    base += ret;
+  }
+
+  return size;
+}
+
+int um_ssl_writev(struct um *machine, VALUE ssl, int argc, VALUE *argv) {
+  size_t total = 0;
+  for (int i = 0; i < argc; i++)
+    total += um_ssl_write_all(machine, ssl, argv[i]);
+  return total;
+}

data/ext/um/um_utils.c

@@ -159,7 +159,7 @@ int um_setup_buffer_ring(struct um *machine, unsigned size, unsigned count) {
   return bg_id;
 }
 
-inline VALUE um_get_string_from_buffer_ring(struct um *machine, int bgid, __s32 result, __u32 flags) {
+inline VALUE um_read_from_buffer_ring(struct um *machine, int bgid, __s32 result, __u32 flags) {
   if (!result) return Qnil;
 
   unsigned buf_idx = flags >> IORING_CQE_BUFFER_SHIFT;