support for Raindrops#size= and Raindrops#evaporate!

This allows limited resizing of the Raindrops memory
area since we always over-allocate due to the required
page aligment for mmap.

It would be nice if mremap() worked with MAP_SHARED,
but it does not and triggers a bus error when attempting
to access the new area.

ref: https://bugzilla.kernel.org/show_bug.cgi?id=8691

3 files changed, 199 insertions, 10 deletions

diff --git a/ext/raindrops/extconf.rb b/ext/raindrops/extconf.rb
index fc2b3fd..08b5f44 100644
--- a/ext/raindrops/extconf.rb
+++ b/ext/raindrops/extconf.rb
@@ -3,6 +3,11 @@ require 'mkmf'
 have_func('mmap', 'sys/mman.h') or abort 'mmap() not found'
 have_func('munmap', 'sys/mman.h') or abort 'munmap() not found'
 
+$CPPFLAGS += " -D_GNU_SOURCE "
+have_func('mremap', 'sys/mman.h')
+
+$CPPFLAGS += " -D_BSD_SOURCE -D_XOPEN_SOURCE=600 "
+have_func("getpagesize", "unistd.h")
 have_func("rb_struct_alloc_noinit")
 have_func('rb_thread_blocking_region')
 
diff --git a/ext/raindrops/raindrops.c b/ext/raindrops/raindrops.c
index dbb076a..1134920 100644
--- a/ext/raindrops/raindrops.c
+++ b/ext/raindrops/raindrops.c
@@ -9,6 +9,9 @@
 #ifndef SIZET2NUM
 #  define SIZET2NUM(x) ULONG2NUM(x)
 #endif
+#ifndef NUM2SIZET
+#  define NUM2SIZET(x) NUM2ULONG(x)
+#endif
 
 /*
  * most modern CPUs have a cache-line size of 64 or 128.
@@ -16,6 +19,10 @@
  * heavily used
  */
 static size_t raindrop_size = 128;
+static size_t rd_page_size;
+
+#define PAGE_MASK               (~(rd_page_size - 1))
+#define PAGE_ALIGN(addr)        (((addr) + rd_page_size - 1) & PAGE_MASK)
 
 /* each raindrop is a counter */
 struct raindrop {
@@ -25,16 +32,18 @@ struct raindrop {
 /* allow mmap-ed regions to store more than one raindrop */
 struct raindrops {
         long size;
+        size_t capa;
+        pid_t pid;
         struct raindrop *drops;
 };
 
 /* called by GC */
-static void evaporate(void *ptr)
+static void gcfree(void *ptr)
 {
         struct raindrops *r = ptr;
 
-        if (r->drops) {
-                int rv = munmap(r->drops, raindrop_size * r->size);
+        if (r->drops != MAP_FAILED) {
+                int rv = munmap(r->drops, raindrop_size * r->capa);
                 if (rv != 0)
                         rb_bug("munmap failed in gc: %s", strerror(errno));
         }
@@ -46,8 +55,10 @@ static void evaporate(void *ptr)
 static VALUE alloc(VALUE klass)
 {
         struct raindrops *r;
+        VALUE rv = Data_Make_Struct(klass, struct raindrops, NULL, gcfree, r);
 
-        return Data_Make_Struct(klass, struct raindrops, NULL, evaporate, r);
+        r->drops = MAP_FAILED;
+        return rv;
 }
 
 static struct raindrops *get(VALUE self)
@@ -56,6 +67,9 @@ static struct raindrops *get(VALUE self)
 
         Data_Get_Struct(self, struct raindrops, r);
 
+        if (r->drops == MAP_FAILED)
+                rb_raise(rb_eStandardError, "invalid or freed Raindrops");
+
         return r;
 }
 
@@ -71,32 +85,118 @@ static struct raindrops *get(VALUE self)
  */
 static VALUE init(VALUE self, VALUE size)
 {
-        struct raindrops *r = get(self);
+        struct raindrops *r = DATA_PTR(self);
         int tries = 1;
+        size_t tmp;
 
-        if (r->drops)
+        if (r->drops != MAP_FAILED)
                 rb_raise(rb_eRuntimeError, "already initialized");
 
         r->size = NUM2LONG(size);
         if (r->size < 1)
                 rb_raise(rb_eArgError, "size must be >= 1");
 
+        tmp = PAGE_ALIGN(raindrop_size * r->size);
+        r->capa = tmp / raindrop_size;
+        assert(PAGE_ALIGN(raindrop_size * r->capa) == tmp && "not aligned");
+
 retry:
-        r->drops = mmap(NULL, raindrop_size * r->size,
+        r->drops = mmap(NULL, tmp,
                         PROT_READ|PROT_WRITE, MAP_ANON|MAP_SHARED, -1, 0);
         if (r->drops == MAP_FAILED) {
-                r->drops = NULL;
                 if ((errno == EAGAIN || errno == ENOMEM) && tries-- > 0) {
                         rb_gc();
                         goto retry;
                 }
                 rb_sys_fail("mmap");
         }
+        r->pid = getpid();
 
         return self;
 }
 
 /*
+ * mremap() is currently broken with MAP_SHARED
+ * https://bugzilla.kernel.org/show_bug.cgi?id=8691
+ */
+#if defined(HAVE_MREMAP) && !defined(MREMAP_WORKS_WITH_MAP_SHARED)
+#  undef HAVE_MREMAP
+#endif
+
+#ifdef HAVE_MREMAP
+#ifndef MREMAP_MAYMOVE
+#  warn MREMAP_MAYMOVE undefined
+#  define MREMAP_MAYMOVE 0
+#endif
+static void resize(struct raindrops *r, size_t new_rd_size)
+{
+        size_t old_size = raindrop_size * r->capa;
+        size_t new_size = PAGE_ALIGN(raindrop_size * new_rd_size);
+        void *old_address = r->drops;
+        void *rv;
+
+        if (r->pid != getpid())
+                rb_raise(rb_eRuntimeError, "cannot mremap() from child");
+
+        rv = mremap(old_address, old_size, new_size, MREMAP_MAYMOVE);
+        if (rv == MAP_FAILED) {
+                if (errno == EAGAIN || errno == ENOMEM) {
+                        rb_gc();
+                        rv = mremap(old_address, old_size, new_size, 0);
+                }
+                if (rv == MAP_FAILED)
+                        rb_sys_fail("mremap");
+        }
+        r->drops = rv;
+        r->size = new_rd_size;
+        r->capa = new_size / raindrop_size;
+        assert(r->capa >= (size_t)r->size && "bad sizing");
+}
+#else /* ! HAVE_MREMAP */
+/*
+ * we cannot use munmap + mmap to reallocate the buffer since it may
+ * already be shared by other processes, so we just fail
+ */
+static void resize(struct raindrops *r, size_t new_capa)
+{
+        rb_raise(rb_eRangeError, "mremap(2) is not available");
+}
+#endif /* ! HAVE_MREMAP */
+
+/*
+ * call-seq:
+ *        rd.size = new_size
+ *
+ * Increases or decreases the current capacity of our Raindrop.
+ * Raises RangeError if +new_size+ is too big or small for the
+ * current backing store
+ */
+static VALUE setsize(VALUE self, VALUE new_size)
+{
+        size_t new_capa = NUM2SIZET(new_size);
+        struct raindrops *r = get(self);
+
+        if (new_capa <= r->capa)
+                r->size = new_capa;
+        else
+                resize(r, new_capa);
+
+        return new_capa;
+}
+
+/*
+ * call-seq:
+ *        rd.capa                -> Integer
+ *
+ * Returns the number of slots allocated (but not necessarily used) by
+ * the Raindrops object.
+ */
+static VALUE capa(VALUE self)
+{
+        return SIZET2NUM(get(self)->capa);
+}
+
+/*
  * call-seq:
  *        rd.dup                -> rd_copy
  *
@@ -104,7 +204,7 @@ retry:
  */
 static VALUE init_copy(VALUE dest, VALUE source)
 {
-        struct raindrops *dst = get(dest);
+        struct raindrops *dst = DATA_PTR(dest);
         struct raindrops *src = get(source);
 
         init(dest, LONG2NUM(src->size));
@@ -234,6 +334,26 @@ void Init_raindrops_linux_tcp_info(void);
 #  endif
 #endif
 
+/*
+ * call-seq:
+ *        rd.evaporate!        -> nil
+ *
+ * Releases mmap()-ed memory allocated for the Raindrops object back
+ * to the OS.  The Ruby garbage collector will also release memory
+ * automatically when it is not needed, but this forces release
+ * under high memory pressure.
+ */
+static VALUE evaporate_bang(VALUE self)
+{
+        struct raindrops *r = get(self);
+        void *addr = r->drops;
+
+        r->drops = MAP_FAILED;
+        if (munmap(addr, raindrop_size * r->capa) != 0)
+                rb_sys_fail("munmap");
+        return Qnil;
+}
+
 void Init_raindrops_ext(void)
 {
         VALUE cRaindrops = rb_define_class("Raindrops", rb_cObject);
@@ -252,6 +372,29 @@ void Init_raindrops_ext(void)
                         raindrop_size = (size_t)tmp;
         }
 #endif
+#if defined(_SC_PAGE_SIZE)
+        rd_page_size = (size_t)sysconf(_SC_PAGE_SIZE);
+#elif defined(_SC_PAGESIZE)
+        rd_page_size = (size_t)sysconf(_SC_PAGESIZE);
+#elif defined(HAVE_GETPAGESIZE)
+        rd_page_size = (size_t)getpagesize();
+#elif defined(PAGE_SIZE)
+        rd_page_size = (size_t)PAGE_SIZE;
+#elif defined(PAGESIZE)
+        rd_page_size = (size_t)PAGESIZE;
+#else
+#  error unable to detect page size for mmap()
+#endif
+        if ((rd_page_size == (size_t)-1) || (rd_page_size < raindrop_size))
+                rb_raise(rb_eRuntimeError,
+                         "system page size invalid: %llu",
+                         (unsigned long long)rd_page_size);
+
+        /*
+         * The size of one page of memory for a mmap()-ed Raindrops region.
+         * Typically 4096 bytes under Linux.
+         */
+        rb_define_const(cRaindrops, "PAGE_SIZE", SIZET2NUM(rd_page_size));
 
         /*
          * The size (in bytes) of a slot in a Raindrops object.
@@ -276,7 +419,10 @@ void Init_raindrops_ext(void)
         rb_define_method(cRaindrops, "[]", aref, 1);
         rb_define_method(cRaindrops, "[]=", aset, 2);
         rb_define_method(cRaindrops, "size", size, 0);
+        rb_define_method(cRaindrops, "size=", setsize, 1);
+        rb_define_method(cRaindrops, "capa", capa, 0);
         rb_define_method(cRaindrops, "initialize_copy", init_copy, 1);
+        rb_define_method(cRaindrops, "evaporate!", evaporate_bang, 0);
 
 #ifdef __linux__
         Init_raindrops_linux_inet_diag();
diff --git a/test/test_raindrops.rb b/test/test_raindrops.rb
index 6686796..a26d0e1 100644
--- a/test/test_raindrops.rb
+++ b/test/test_raindrops.rb
@@ -16,9 +16,15 @@ class TestRaindrops < Test::Unit::TestCase
     puts "Raindrops::SIZE = #{Raindrops::SIZE}"
   end
 
-  def test_size
+  def test_page_size
+    assert_kind_of Integer, Raindrops::PAGE_SIZE
+    assert Raindrops::PAGE_SIZE > Raindrops::SIZE
+  end
+
+  def test_size_and_capa
     rd = Raindrops.new(4)
     assert_equal 4, rd.size
+    assert rd.capa >= rd.size
   end
 
   def test_ary
@@ -104,4 +110,36 @@ class TestRaindrops < Test::Unit::TestCase
     assert_equal expect, rd.to_ary
   end
 
+  def test_resize
+    rd = Raindrops.new(4)
+    assert_equal 4, rd.size
+    assert_equal rd.capa, rd.size = rd.capa
+    assert_equal rd.capa, rd.to_ary.size
+    assert_equal 0, rd[rd.capa - 1]
+    assert_equal 1, rd.incr(rd.capa - 1)
+    assert_raises(ArgumentError) { rd[rd.capa] }
+  end
+
+  def test_resize_mremap
+    rd = Raindrops.new(4)
+    assert_equal 4, rd.size
+    old_capa = rd.capa
+    rd.size = rd.capa + 1
+    assert_equal old_capa * 2, rd.capa
+
+    # mremap() is currently broken with MAP_SHARED
+    # https://bugzilla.kernel.org/show_bug.cgi?id=8691
+    assert_equal 0, rd[old_capa]
+    assert_equal rd.capa, rd.to_ary.size
+    assert_equal 0, rd[rd.capa - 1]
+    assert_equal 1, rd.incr(rd.capa - 1)
+    assert_raises(ArgumentError) { rd[rd.capa] }
+    rescue RangeError
+  end # if RUBY_PLATFORM =~ /linux/
+
+  def test_evaporate
+    rd = Raindrops.new 1
+    assert_nil rd.evaporate!
+    assert_raises(StandardError) { rd.evaporate! }
+  end
 end