vibe_zstd 1.1.0 → 1.3.0

data/ext/vibe_zstd/streaming.c

@@ -14,6 +14,12 @@ static VALUE vibe_zstd_reader_initialize(int argc, VALUE *argv, VALUE self);
 static VALUE vibe_zstd_reader_read(int argc, VALUE *argv, VALUE self);
 static VALUE vibe_zstd_reader_eof(VALUE self);
 
+// State struct for rb_ensure-based string lock/unlock in vibe_zstd_writer_write
+typedef struct {
+    vibe_zstd_cstream* cstream;
+    VALUE data;
+} vibe_zstd_write_state;
+
 // TypedData types - defined in vibe_zstd.c
 extern rb_data_type_t vibe_zstd_cstream_type;
 extern rb_data_type_t vibe_zstd_dstream_type;
@@ -89,6 +95,9 @@ vibe_zstd_writer_initialize(int argc, VALUE *argv, VALUE self) {
         if (ZSTD_isError(result)) {
             rb_raise(rb_eRuntimeError, "Failed to set dictionary: %s", ZSTD_getErrorName(result));
         }
+        // Retain the CDict object so GC won't free it while the stream holds a raw
+        // pointer to its internal ZSTD_CDict (ZSTD_CCtx_refCDict stores no Ruby ref)
+        rb_ivar_set(self, rb_intern("@dict"), dict);
     }
 
     // Allocate reusable output buffer (write barrier for WB_PROTECTED)
@@ -97,14 +106,16 @@ vibe_zstd_writer_initialize(int argc, VALUE *argv, VALUE self) {
     return self;
 }
 
+// Body of the rb_ensure wrapper: runs the compress loop with data locked
 static VALUE
-vibe_zstd_writer_write(VALUE self, VALUE data) {
-    Check_Type(data, T_STRING);
-
-    vibe_zstd_cstream* cstream;
-    TypedData_Get_Struct(self, vibe_zstd_cstream, &vibe_zstd_cstream_type, cstream);
-
-    // Input buffer: pos advances as ZSTD consumes data
+vibe_zstd_writer_write_body(VALUE arg) {
+    vibe_zstd_write_state* state = (vibe_zstd_write_state*)arg;
+    vibe_zstd_cstream* cstream = state->cstream;
+    VALUE data = state->data;
+
+    // Input buffer: pos advances as ZSTD consumes data.
+    // data is locked (rb_str_locktmp) for the duration of this call so that
+    // RSTRING_PTR remains valid even when rb_funcall runs arbitrary Ruby code.
     ZSTD_inBuffer input = {
         .src = RSTRING_PTR(data),
         .size = RSTRING_LEN(data),
@@ -116,7 +127,11 @@ vibe_zstd_writer_write(VALUE self, VALUE data) {
 
     // Process all input data in chunks
     while (input.pos < input.size) {
-        rb_str_set_len(outBuffer, 0);  // Reset buffer for reuse
+        // Unshare buffer if COW-shared by a prior IO#write receiver (Ruby 3.3+),
+        // then restore capacity which may have shrunk during unsharing
+        rb_str_modify(outBuffer);
+        rb_str_resize(outBuffer, (long)outBufferSize);
+        rb_str_set_len(outBuffer, 0);
         ZSTD_outBuffer output = {
             .dst = RSTRING_PTR(outBuffer),
             .size = outBufferSize,
@@ -133,10 +148,39 @@ vibe_zstd_writer_write(VALUE self, VALUE data) {
         // Write any compressed output that was produced
         if (output.pos > 0) {
             rb_str_set_len(outBuffer, output.pos);
+            // rb_funcall may run arbitrary Ruby code, but input.src stays valid
+            // because data is locked against mutation/reallocation
             rb_funcall(cstream->io, id_write, 1, outBuffer);
         }
     }
 
+    return Qnil;
+}
+
+// Ensure function: always unlocks data regardless of raise/return
+static VALUE
+vibe_zstd_writer_write_unlock(VALUE arg) {
+    rb_str_unlocktmp((VALUE)arg);
+    return Qnil;
+}
+
+static VALUE
+vibe_zstd_writer_write(VALUE self, VALUE data) {
+    Check_Type(data, T_STRING);
+
+    vibe_zstd_cstream* cstream;
+    TypedData_Get_Struct(self, vibe_zstd_cstream, &vibe_zstd_cstream_type, cstream);
+
+    // Lock data for the duration of the compress loop so that RSTRING_PTR(data)
+    // stays valid even when io.write (called inside the loop) runs Ruby code that
+    // could otherwise mutate or resize the string.  rb_str_locktmp raises if the
+    // string is already locked; the ensure always unlocks it.
+    rb_str_locktmp(data);
+
+    vibe_zstd_write_state state = { cstream, data };
+    rb_ensure(vibe_zstd_writer_write_body, (VALUE)&state,
+              vibe_zstd_writer_write_unlock, data);
+
     return self;
 }
 
@@ -154,7 +198,9 @@ vibe_zstd_writer_flush(VALUE self) {
     // ZSTD_e_flush: flush internal buffers, making all data readable
     // Loop until remaining == 0 (flush complete)
     do {
-        rb_str_set_len(outBuffer, 0);  // Reset buffer for reuse
+        rb_str_modify(outBuffer);
+        rb_str_resize(outBuffer, (long)outBufferSize);
+        rb_str_set_len(outBuffer, 0);
         ZSTD_outBuffer output = {
             .dst = RSTRING_PTR(outBuffer),
             .size = outBufferSize,
@@ -190,7 +236,9 @@ vibe_zstd_writer_finish(VALUE self) {
     // ZSTD_e_end: finalize frame with checksum and epilogue
     // Loop until remaining == 0 (frame complete)
     do {
-        rb_str_set_len(outBuffer, 0);  // Reset buffer for reuse
+        rb_str_modify(outBuffer);
+        rb_str_resize(outBuffer, (long)outBufferSize);
+        rb_str_set_len(outBuffer, 0);
         ZSTD_outBuffer output = {
             .dst = RSTRING_PTR(outBuffer),
             .size = outBufferSize,
@@ -267,6 +315,9 @@ vibe_zstd_reader_initialize(int argc, VALUE *argv, VALUE self) {
         if (ZSTD_isError(result)) {
             rb_raise(rb_eRuntimeError, "Failed to set dictionary: %s", ZSTD_getErrorName(result));
         }
+        // Retain the DDict object so GC won't free it while the stream holds a raw
+        // pointer to its internal ZSTD_DDict (ZSTD_DCtx_refDDict stores no Ruby ref)
+        rb_ivar_set(self, rb_intern("@dict"), dict);
     }
 
     // Initialize input buffer management
@@ -290,10 +341,18 @@ vibe_zstd_reader_initialize(int argc, VALUE *argv, VALUE self) {
 // - Maintains internal compressed input buffer that refills from IO as needed
 // - Calls ZSTD_decompressStream incrementally to produce output
 // - Tracks EOF state based on IO exhaustion and frame completion
+// - Input chunks are stored as frozen copies so that IOs which mutate/reuse
+//   the returned string cannot invalidate dstream->input.src between calls
 //
 // EOF handling:
 // - Returns nil when no more data available
 // - Sets eof flag when: IO returns nil, frame complete (ret==0), or no progress made
+// - read(0) always returns "" immediately without touching stream state
+//
+// Allocation strategy:
+// - Initial buffer is capped at ZSTD_DStreamOutSize() to avoid gigabyte
+//   allocations for large size arguments on small streams
+// - Buffer grows geometrically (doubling) up to requested_size as needed
 //
 // This implements proper streaming semantics for incremental decompression
 // of arbitrarily large files without loading everything into memory.
@@ -305,6 +364,11 @@ vibe_zstd_reader_read(int argc, VALUE *argv, VALUE self) {
     vibe_zstd_dstream* dstream;
     TypedData_Get_Struct(self, vibe_zstd_dstream, &vibe_zstd_dstream_type, dstream);
 
+    // read(0): per IO semantics, always return "" without touching stream state
+    if (!NIL_P(size_arg) && NUM2SIZET(size_arg) == 0) {
+        return rb_str_new(NULL, 0);
+    }
+
     if (dstream->eof) {
         return Qnil;
     }
@@ -315,8 +379,12 @@ vibe_zstd_reader_read(int argc, VALUE *argv, VALUE self) {
     size_t requested_size = NIL_P(size_arg) ? default_chunk_size : NUM2SIZET(size_arg);
     size_t inBufferSize = ZSTD_DStreamInSize();
 
-    // Preallocate buffer for requested size
-    VALUE result = rb_str_buf_new(requested_size);
+    // Cap the initial allocation to avoid multi-gigabyte pre-allocations when
+    // the caller passes a huge size argument for a small stream.  The buffer
+    // grows geometrically below as output accumulates.
+    size_t default_out_size = ZSTD_DStreamOutSize();
+    size_t initial_alloc = (requested_size < default_out_size) ? requested_size : default_out_size;
+    VALUE result = rb_str_buf_new((long)initial_alloc);
 
     size_t total_read = 0;
     int made_progress = 0;
@@ -333,10 +401,21 @@ vibe_zstd_reader_read(int argc, VALUE *argv, VALUE self) {
                 break;
             }
 
+            // The IO is duck-typed: read may return anything. Convert via to_str
+            // (raising TypeError otherwise) so RSTRING below never sees a non-String.
+            StringValue(chunk);
+
+            // Store a private frozen copy so that an IO that reuses/mutates its
+            // returned buffer string cannot invalidate dstream->input.src between
+            // successive read() calls.  rb_str_new_frozen is cheap (copy-on-write
+            // snapshot) when the string is already frozen, and allocates a
+            // separate copy otherwise.
+            VALUE frozen_chunk = rb_str_new_frozen(chunk);
+
             // Reset input buffer with new data (write barrier for WB_PROTECTED)
-            RB_OBJ_WRITE(self, &dstream->input_data, chunk);
-            dstream->input.src = RSTRING_PTR(chunk);
-            dstream->input.size = RSTRING_LEN(chunk);
+            RB_OBJ_WRITE(self, &dstream->input_data, frozen_chunk);
+            dstream->input.src = RSTRING_PTR(frozen_chunk);
+            dstream->input.size = RSTRING_LEN(frozen_chunk);
             dstream->input.pos = 0;
         }
 
@@ -345,7 +424,22 @@ vibe_zstd_reader_read(int argc, VALUE *argv, VALUE self) {
             break;
         }
 
-        size_t space_left = requested_size - total_read;
+        // Grow the output buffer geometrically when it is full, capped at
+        // requested_size.  We must recompute RSTRING_PTR after any resize
+        // because the backing allocation may move.
+        size_t current_capacity = (size_t)rb_str_capacity(result);
+        if (total_read >= current_capacity) {
+            size_t new_capacity = current_capacity * 2;
+            if (new_capacity > requested_size) new_capacity = requested_size;
+            rb_str_resize(result, (long)new_capacity);
+        }
+
+        // Cap space_left at (requested_size - total_read) to ensure read(n) never
+        // returns more than n bytes: rb_str_capacity may exceed the requested size
+        // due to malloc's internal size-class rounding (e.g. request 100, get 135).
+        size_t effective_capacity = (size_t)rb_str_capacity(result);
+        if (effective_capacity > requested_size) effective_capacity = requested_size;
+        size_t space_left = effective_capacity - total_read;
 
         ZSTD_outBuffer output = {
             .dst = RSTRING_PTR(result) + total_read,

data/ext/vibe_zstd/vibe_zstd.c

@@ -210,6 +210,7 @@ vibe_zstd_dctx_alloc(VALUE klass) {
         rb_raise(rb_eRuntimeError, "Failed to create ZSTD_DCtx");
     }
     dctx->initial_capacity = 0;  // 0 = use class default
+    dctx->max_decompressed_size = 0;  // 0 = inherit class default
     return TypedData_Wrap_Struct(klass, &vibe_zstd_dctx_type, dctx);
 }
 
@@ -279,6 +280,35 @@ vibe_zstd_default_c_level(VALUE self) {
     return INT2NUM(ZSTD_defaultCLevel());
 }
 
+// Run func(arg) without the GVL while str is locked against mutation.
+// rb_thread_call_without_gvl can deliver a pending async exception (e.g.
+// Thread#raise, Timeout) after reacquiring the GVL, so the unlock must go
+// through rb_ensure or the string would be left permanently locked.
+typedef struct {
+    void* (*func)(void*);
+    void* arg;
+} nogvl_locked_call;
+
+static VALUE
+nogvl_locked_body(VALUE p) {
+    nogvl_locked_call* call = (nogvl_locked_call*)p;
+    rb_thread_call_without_gvl(call->func, call->arg, NULL, NULL);
+    return Qnil;
+}
+
+static VALUE
+nogvl_locked_unlock(VALUE str) {
+    rb_str_unlocktmp(str);
+    return Qnil;
+}
+
+static void
+vibe_zstd_nogvl_with_str_locked(void* (*func)(void*), void* arg, VALUE str) {
+    nogvl_locked_call call = { func, arg };
+    rb_str_locktmp(str);
+    rb_ensure(nogvl_locked_body, (VALUE)&call, nogvl_locked_unlock, str);
+}
+
 // Include the split implementation files
 #include "cctx.c"
 #include "dctx.c"

data/ext/vibe_zstd/vibe_zstd.h

@@ -13,6 +13,7 @@ typedef struct {
 typedef struct {
     ZSTD_DCtx* dctx;
     size_t initial_capacity;  // Initial capacity for unknown-size decompression (0 = use class default)
+    size_t max_decompressed_size;  // Output size limit (0 = inherit class default; class default 0 = unlimited)
 } vibe_zstd_dctx;
 
 typedef struct {

data/lib/vibe_zstd/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module VibeZstd
-  VERSION = "1.1.0"
+  VERSION = "1.3.0"
 end

data/lib/vibe_zstd.rb

@@ -7,18 +7,33 @@ require_relative "vibe_zstd/constants"
 module VibeZstd
   class Error < StandardError; end
 
-  # Convenience method for one-off compression
-  # Supports all CCtx#compress options: level, dict, pledged_size
+  # Keyword options handled per-operation by CCtx#compress / DCtx#decompress.
+  # Any other keyword is treated as a context (sticky) parameter and applied via
+  # the constructor, so e.g. VibeZstd.compress(data, checksum_flag: true) and
+  # VibeZstd.decompress(data, format: 1) work through the convenience methods.
+  # An unknown keyword raises NoMethodError from the corresponding setter.
+  COMPRESS_CALL_OPTIONS = %i[level dict pledged_size].freeze
+  DECOMPRESS_CALL_OPTIONS = %i[dict initial_capacity max_decompressed_size max_size].freeze
+  private_constant :COMPRESS_CALL_OPTIONS, :DECOMPRESS_CALL_OPTIONS
+
+  # Convenience method for one-off compression.
+  # Per-call options (level, dict, pledged_size) are passed to #compress; any
+  # other keyword is a context parameter (e.g. checksum_flag:, window_log:,
+  # workers:, format:) applied to a fresh CCtx.
   def self.compress(data, **options)
-    cctx = CCtx.new
-    cctx.compress(data, **options)
+    call_opts = options.slice(*COMPRESS_CALL_OPTIONS)
+    ctx_opts = options.except(*COMPRESS_CALL_OPTIONS)
+    CCtx.new(**ctx_opts).compress(data, **call_opts)
   end
 
-  # Convenience method for one-off decompression
-  # Supports all DCtx#decompress options: dict, initial_capacity
+  # Convenience method for one-off decompression.
+  # Per-call options (dict, initial_capacity, max_decompressed_size/max_size) are
+  # passed to #decompress; any other keyword is a context parameter (e.g.
+  # format:, window_log_max:) applied to a fresh DCtx.
   def self.decompress(data, **options)
-    dctx = DCtx.new
-    dctx.decompress(data, **options)
+    call_opts = options.slice(*DECOMPRESS_CALL_OPTIONS)
+    ctx_opts = options.except(*DECOMPRESS_CALL_OPTIONS)
+    DCtx.new(**ctx_opts).decompress(data, **call_opts)
   end
 
   # Get the decompressed content size from a compressed frame
@@ -88,6 +103,10 @@ module VibeZstd
   # Memory footprint: ~128KB per DCtx × unique dictionaries × threads
   # Example: 3 dicts × 5 Puma threads = 1.9MB total
   #
+  # Storage: uses Thread#thread_variable_get/set (true thread-local) so that
+  # fiber-based servers (Falcon, async) share one pool per OS thread rather
+  # than allocating a fresh pool for every fiber.
+  #
   # Note: Only supports per-operation parameters (level, dict, pledged_size, initial_capacity)
   # Does NOT support context-level settings (nb_workers, checksum_flag, etc.)
   module ThreadLocal
@@ -103,9 +122,10 @@ module VibeZstd
       # Key by dictionary ID, or :default if no dict
       key = dict ? dict.dict_id : :default
 
-      # Get or create thread-local context pool
-      Thread.current[:vibe_zstd_cctx_pool] ||= {}
-      cctx = Thread.current[:vibe_zstd_cctx_pool][key] ||= VibeZstd::CCtx.new
+      # Get or create thread-local context pool (true thread-local, not fiber-local)
+      pool = Thread.current.thread_variable_get(:vibe_zstd_cctx_pool) || {}
+      cctx = pool[key] ||= VibeZstd::CCtx.new
+      Thread.current.thread_variable_set(:vibe_zstd_cctx_pool, pool)
 
       # Build options hash
       options = {}
@@ -122,18 +142,21 @@ module VibeZstd
     # @param data [String] Data to decompress
     # @param dict [DDict] Decompression dictionary (optional)
     # @param initial_capacity [Integer] Initial buffer size for unknown-size frames (optional)
+    # @param max_decompressed_size [Integer] Output-size limit; raises DecompressedSizeExceeded if exceeded (optional)
     # @return [String] Decompressed data
-    def self.decompress(data, dict: nil, initial_capacity: nil)
+    def self.decompress(data, dict: nil, initial_capacity: nil, max_decompressed_size: nil)
       key = dict ? dict.dict_id : :default
 
-      # Get or create thread-local context pool
-      Thread.current[:vibe_zstd_dctx_pool] ||= {}
-      dctx = Thread.current[:vibe_zstd_dctx_pool][key] ||= VibeZstd::DCtx.new
+      # Get or create thread-local context pool (true thread-local, not fiber-local)
+      pool = Thread.current.thread_variable_get(:vibe_zstd_dctx_pool) || {}
+      dctx = pool[key] ||= VibeZstd::DCtx.new
+      Thread.current.thread_variable_set(:vibe_zstd_dctx_pool, pool)
 
       # Build options hash
       options = {}
       options[:dict] = dict if dict
       options[:initial_capacity] = initial_capacity if initial_capacity
+      options[:max_decompressed_size] = max_decompressed_size if max_decompressed_size
 
       # C code will validate dict matches frame requirements
       dctx.decompress(data, **options)
@@ -142,19 +165,21 @@ module VibeZstd
     # Clear all thread-local context pools for the current thread
     # Useful for testing or explicit memory management
     def self.clear_thread_cache!
-      Thread.current[:vibe_zstd_cctx_pool] = {}
-      Thread.current[:vibe_zstd_dctx_pool] = {}
+      Thread.current.thread_variable_set(:vibe_zstd_cctx_pool, {})
+      Thread.current.thread_variable_set(:vibe_zstd_dctx_pool, {})
       nil
     end
 
     # Get statistics about the current thread's context pools
     # @return [Hash] Pool statistics
     def self.thread_cache_stats
+      cctx_pool = Thread.current.thread_variable_get(:vibe_zstd_cctx_pool)
+      dctx_pool = Thread.current.thread_variable_get(:vibe_zstd_dctx_pool)
       {
-        compression_contexts: Thread.current[:vibe_zstd_cctx_pool]&.size || 0,
-        decompression_contexts: Thread.current[:vibe_zstd_dctx_pool]&.size || 0,
-        compression_keys: Thread.current[:vibe_zstd_cctx_pool]&.keys || [],
-        decompression_keys: Thread.current[:vibe_zstd_dctx_pool]&.keys || []
+        compression_contexts: cctx_pool&.size || 0,
+        decompression_contexts: dctx_pool&.size || 0,
+        compression_keys: cctx_pool&.keys || [],
+        decompression_keys: dctx_pool&.keys || []
       }
     end
   end
@@ -229,7 +254,7 @@ module VibeZstd
       loop do
         # Check buffer for separator
         if (idx = @line_buffer.index(sep))
-          return @line_buffer.slice!(0, idx + sep.bytesize)
+          return @line_buffer.slice!(0, idx + sep.length)
         end
 
         # Read more data in larger chunks
@@ -240,7 +265,7 @@ module VibeZstd
       end
 
       # Return remaining buffer or nil
-      @line_buffer.empty? ? nil : @line_buffer.slice!(0, @line_buffer.bytesize)
+      @line_buffer.empty? ? nil : @line_buffer.slice!(0, @line_buffer.length)
     end
 
     # Iterate over lines

metadata

@@ -1,13 +1,13 @@
 --- !ruby/object:Gem::Specification
 name: vibe_zstd
 version: !ruby/object:Gem::Version
-  version: 1.1.0
+  version: 1.3.0
 platform: ruby
 authors:
 - Kelley Reynolds
 bindir: exe
 cert_chain: []
-date: 2026-03-03 00:00:00.000000000 Z
+date: 2026-06-12 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: benchmark-ips
@@ -64,6 +64,7 @@ files:
 - benchmark/streaming.rb
 - ext/vibe_zstd/cctx.c
 - ext/vibe_zstd/dctx.c
+- ext/vibe_zstd/depend
 - ext/vibe_zstd/dict.c
 - ext/vibe_zstd/extconf.rb
 - ext/vibe_zstd/frames.c