add Socket and Thread.close!

Author: Asmod4n

mrblib/socket.rb

@@ -5,8 +5,14 @@ def bind(endpoint)
       self
     end
 
-    def close(blocky = true)
-      LibZMQ.close(self, blocky)
+    def close
+      LibZMQ.close(self)
+      nil
+    end
+
+    # this immediately closes the socket, discarding any messages which couldn't be sent yet.
+    def close!
+      LibZMQ.close!(self)
       nil
     end
 

mrblib/thread.rb

@@ -68,8 +68,12 @@ def finalize(object_id)
       LibZMQ.send(@pipe, [FINALIZE, object_id].to_msgpack, 0)
     end
 
-    def close(blocky = true)
-      LibZMQ.threadclose(self, blocky)
+    def close
+      LibZMQ.threadclose(self)
+    end
+
+    def close!
+      LibZMQ.threadclose!(self)
     end
 
     class Thread_fn

src/mrb_libzmq.c

@@ -19,14 +19,28 @@ static mrb_value
 mrb_zmq_close(mrb_state *mrb, mrb_value self)
 {
   mrb_value socket_val;
-  mrb_bool blocky = TRUE;
-  mrb_get_args(mrb, "o|b", &socket_val, &blocky);
+  mrb_get_args(mrb, "o", &socket_val);
 
   if (mrb_type(socket_val) == MRB_TT_DATA && DATA_TYPE(socket_val) == &mrb_zmq_socket_type) {
-    if (!blocky) {
-      int linger = 0;
-      zmq_setsockopt(DATA_PTR(socket_val), ZMQ_LINGER, &linger, sizeof(linger));
+    int rc = zmq_close(DATA_PTR(socket_val));
+    if (unlikely(rc == -1)) {
+      mrb_zmq_handle_error(mrb, "zmq_close");
     }
+    mrb_data_init(socket_val, NULL, NULL);
+  }
+
+  return mrb_nil_value();
+}
+
+static mrb_value
+mrb_zmq_close_mark(mrb_state *mrb, mrb_value self)
+{
+  mrb_value socket_val;
+  mrb_get_args(mrb, "o", &socket_val);
+
+  if (mrb_type(socket_val) == MRB_TT_DATA && DATA_TYPE(socket_val) == &mrb_zmq_socket_type) {
+    int disable = 0;
+    zmq_setsockopt(DATA_PTR(socket_val), ZMQ_LINGER, &disable, sizeof(disable));
     int rc = zmq_close(DATA_PTR(socket_val));
     if (unlikely(rc == -1)) {
       mrb_zmq_handle_error(mrb, "zmq_close");
@@ -516,7 +530,7 @@ mrb_zmq_socket(mrb_state *mrb, mrb_value self)
   mrb_get_args(mrb, "i", &type);
   mrb_assert_int_fit(mrb_int, type, int, INT_MAX);
 
-  void *socket = zmq_socket(MRB_LIBZMQ_CONTEXT(mrb), type);
+  void *socket = zmq_socket(MRB_LIBZMQ_CONTEXT(mrb), (int) type);
   if (likely(socket)) {
     mrb_data_init(self, socket, &mrb_zmq_socket_type);
   } else {
@@ -535,7 +549,7 @@ mrb_zmq_socket_monitor(mrb_state *mrb, mrb_value self)
   mrb_get_args(mrb, "dzi", &socket, &mrb_zmq_socket_type, &addr, &events);
   mrb_assert_int_fit(mrb_int, events, int, INT_MAX);
 
-  int rc = zmq_socket_monitor(socket, addr, events);
+  int rc = zmq_socket_monitor(socket, addr, (int) events);
   if (unlikely(rc == -1)) {
     mrb_zmq_handle_error(mrb, "zmq_socket_monitor");
   }
@@ -558,7 +572,7 @@ mrb_zmq_socket_recv(mrb_state *mrb, mrb_value self)
     mrb_value msg_val = mrb_obj_new(mrb, zmq_msg_class, 0, NULL);
     zmq_msg_t *msg = (zmq_msg_t *) DATA_PTR(msg_val);
 
-    int rc = zmq_msg_recv (msg, DATA_PTR(self), flags);
+    int rc = zmq_msg_recv (msg, DATA_PTR(self), (int) flags);
     if (unlikely(rc == -1)) {
       mrb_zmq_handle_error(mrb, "zmq_msg_recv");
     }
@@ -629,11 +643,11 @@ mrb_zmq_thread_fn_cb(mrb_state *mrb, const mrb_value mrb_zmq_thread_data_)
   mrb_zmq_thread_data_t *mrb_zmq_thread_data = (mrb_zmq_thread_data_t *) mrb_cptr(mrb_zmq_thread_data_);
 
   struct RClass* zmq_mod = mrb_module_get(mrb, "ZMQ");
-  mrb_value pipe_val = mrb_obj_value(mrb_obj_alloc(mrb, MRB_TT_DATA, mrb_class_get_under(mrb, zmq_mod, "Pair")));
-  mrb_data_init(pipe_val, mrb_zmq_thread_data->backend, &mrb_zmq_socket_type);
+  mrb_value backend_pipe = mrb_obj_value(mrb_obj_alloc(mrb, MRB_TT_DATA, mrb_class_get_under(mrb, zmq_mod, "Pair")));
+  mrb_data_init(backend_pipe, mrb_zmq_thread_data->backend, &mrb_zmq_socket_type);
   mrb_value timeo = mrb_int_value(mrb, 120000);
-  mrb_funcall(mrb, pipe_val, "sndtimeo=", 1, timeo);
-  mrb_funcall(mrb, pipe_val, "rcvtimeo=", 1, timeo);
+  mrb_funcall(mrb, backend_pipe, "sndtimeo=", 1, timeo);
+  mrb_funcall(mrb, backend_pipe, "rcvtimeo=", 1, timeo);
 
   mrb_value argv = mrb_msgpack_unpack(mrb, mrb_zmq_thread_data->argv_str);
   mrb_value block = mrb_msgpack_unpack(mrb, mrb_zmq_thread_data->block_str);
@@ -645,7 +659,7 @@ mrb_zmq_thread_fn_cb(mrb_state *mrb, const mrb_value mrb_zmq_thread_data_)
   } else {
     mrb_zmq_thread_data->thread_fn = mrb_obj_value(mrb_obj_alloc(mrb, MRB_TT_OBJECT, mrb_class_get_under(mrb, mrb_class_get_under(mrb, zmq_mod, "Thread"), "Thread_fn")));
   }
-  mrb_iv_set(mrb, mrb_zmq_thread_data->thread_fn, mrb_intern_lit(mrb, "@pipe"), pipe_val);
+  mrb_iv_set(mrb, mrb_zmq_thread_data->thread_fn, mrb_intern_lit(mrb, "@pipe"), backend_pipe);
   mrb_funcall_with_block(mrb, mrb_zmq_thread_data->thread_fn, mrb_intern_lit(mrb, "initialize"), RARRAY_LEN(argv), RARRAY_PTR(argv), block);
   int success = TRUE;
   int rc = zmq_send(mrb_zmq_thread_data->backend, &success, sizeof(success), 0);
@@ -731,12 +745,12 @@ mrb_zmq_threadstart(mrb_state *mrb, mrb_value thread_class)
     mrb_true_value()
   };
   mrb_zmq_thread_data->endpoint = mrb_string_value_cstr(mrb, &pipe_args[0]);
-  mrb_value frontend_val = mrb_obj_new(mrb, mrb_class_get_under(mrb, mrb_module_get(mrb, "ZMQ"), "Pair"), NELEMS(pipe_args), pipe_args);
-  mrb_zmq_thread_data->frontend = DATA_PTR(frontend_val);
+  mrb_value frontend_pipe = mrb_obj_new(mrb, mrb_class_get_under(mrb, mrb_module_get(mrb, "ZMQ"), "Pair"), NELEMS(pipe_args), pipe_args);
+  mrb_zmq_thread_data->frontend = DATA_PTR(frontend_pipe);
   mrb_value timeo = mrb_int_value(mrb, 120000);
-  mrb_funcall(mrb, frontend_val, "sndtimeo=", 1, timeo);
-  mrb_funcall(mrb, frontend_val, "rcvtimeo=", 1, timeo);
-  mrb_iv_set(mrb, self, mrb_intern_lit(mrb, "@pipe"), frontend_val);
+  mrb_funcall(mrb, frontend_pipe, "sndtimeo=", 1, timeo);
+  mrb_funcall(mrb, frontend_pipe, "rcvtimeo=", 1, timeo);
+  mrb_iv_set(mrb, self, mrb_intern_lit(mrb, "@pipe"), frontend_pipe);
   mrb_zmq_thread_data->backend = zmq_socket(MRB_LIBZMQ_CONTEXT(mrb), ZMQ_PAIR);
   if (unlikely(!mrb_zmq_thread_data->backend)) {
     mrb_zmq_handle_error(mrb, "zmq_socket");
@@ -751,10 +765,9 @@ mrb_zmq_threadstart(mrb_state *mrb, mrb_value thread_class)
   int success = FALSE;
   zmq_recv(mrb_zmq_thread_data->frontend, &success, sizeof(success), 0);
   if (unlikely(!success)) {
-    mrb_bool rcvmore = mrb_bool(mrb_funcall(mrb, frontend_val, "rcvmore?", 0));
+    mrb_bool rcvmore = mrb_bool(mrb_funcall(mrb, frontend_pipe, "rcvmore?", 0));
     if (rcvmore) {
-      mrb_value exc = mrb_msgpack_unpack(mrb, mrb_funcall(mrb, mrb_funcall(mrb, frontend_val, "recv", 0), "to_str", 0));
-      mrb_exc_raise(mrb, exc);
+      mrb_exc_raise(mrb, mrb_msgpack_unpack(mrb, mrb_funcall(mrb, mrb_funcall(mrb, frontend_pipe, "recv", 0), "to_str", 0)));
     } else {
       mrb_raise(mrb, E_RUNTIME_ERROR, "Cannot initialize ZMQ Thread");
     }
@@ -996,7 +1009,9 @@ mrb_zmq_poller_wait(mrb_state *mrb, mrb_value self)
     return self;
   }
 }
+
 #else
+
 static mrb_value
 mrb_zmq_poller_new(mrb_state *mrb, mrb_value self)
 {
@@ -1108,10 +1123,7 @@ mrb_zmq_poller_remove(mrb_state *mrb, mrb_value self)
           ++ptr;
         }
         mrb_funcall(mrb, sockets, "delete_at", 1, mrb_int_value(mrb, i));
-        pollitems = (struct zmq_pollitem_t *) mrb_realloc(mrb, DATA_PTR(self), RARRAY_LEN(sockets) * sizeof(*pollitems));
-        if (DATA_PTR(self) != pollitems) {
-          mrb_data_init(self, pollitems, &mrb_zmq_poller_type);
-        }
+        DATA_PTR(self) = (struct zmq_pollitem_t *) mrb_realloc(mrb, DATA_PTR(self), RARRAY_LEN(sockets) * sizeof(*pollitems));
 
         return self;
       }
@@ -1216,7 +1228,7 @@ mrb_zmq_timers_add(mrb_state *mrb, mrb_value self)
   timer_fn_arg->block = block;
   mrb_iv_set(mrb, timer, mrb_intern_lit(mrb, "block"), block);
 
-  int timer_id = zmq_timers_add(DATA_PTR(self), interval, mrb_zmq_timer_fn, timer_fn_arg);
+  int timer_id = zmq_timers_add(DATA_PTR(self), (size_t) interval, mrb_zmq_timer_fn, timer_fn_arg);
   if (unlikely(timer_id == -1)) {
     mrb_zmq_handle_error(mrb, "zmq_timers_add");
   }
@@ -1234,7 +1246,7 @@ mrb_zmq_timers_set_interval(mrb_state *mrb, mrb_value self)
   mrb_assert_int_fit(mrb_int, interval, size_t, SIZE_MAX);
 
   mrb_zmq_timers_fn_t *timer_fn_arg = (mrb_zmq_timers_fn_t *) DATA_PTR(self);
-  int rc = zmq_timers_set_interval(DATA_PTR(timer_fn_arg->timers), timer_fn_arg->timer_id, interval);
+  int rc = zmq_timers_set_interval(DATA_PTR(timer_fn_arg->timers), timer_fn_arg->timer_id, (size_t) interval);
   if (unlikely(rc == -1)) {
     mrb_zmq_handle_error(mrb, "zmq_timers_set_interval");
   }
@@ -1426,7 +1438,8 @@ mrb_mruby_zmq_gem_init(mrb_state* mrb)
   libzmq_mod = mrb_define_module(mrb, "LibZMQ");
   mrb_define_const(mrb, libzmq_mod, "_Context", mrb_cptr_value(mrb, context));
   mrb_define_module_function(mrb, libzmq_mod, "bind", mrb_zmq_bind, MRB_ARGS_REQ(2));
-  mrb_define_module_function(mrb, libzmq_mod, "close", mrb_zmq_close, MRB_ARGS_ARG(1, 1));
+  mrb_define_module_function(mrb, libzmq_mod, "close", mrb_zmq_close, MRB_ARGS_REQ(1));
+  mrb_define_module_function(mrb, libzmq_mod, "close!", mrb_zmq_close_mark, MRB_ARGS_REQ(1));
   mrb_define_module_function(mrb, libzmq_mod, "connect", mrb_zmq_connect, MRB_ARGS_REQ(2));
   mrb_define_module_function(mrb, libzmq_mod, "ctx_get", mrb_zmq_ctx_get, MRB_ARGS_REQ(1));
   mrb_define_module_function(mrb, libzmq_mod, "ctx_set", mrb_zmq_ctx_set, MRB_ARGS_REQ(2));