pg-replication-protocol 0.0.7 → 0.0.8

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: fc030b9fec7b0a01acbf309813046913b36da4ddc289e7201d17e78b6de0d471
-  data.tar.gz: 7fa53e2ce183c315b70812f1d58af99afa492419d208704f2cbf97d67a7e57f8
+  metadata.gz: f00c769aee5ae5e1d41063e827f2d5fcaa19991aed6d4676f414ce5f6d6326ef
+  data.tar.gz: a99c6f062f9c33109a409b58ef465833fb1d0b3fd8499a91461db9ea3e2586af
 SHA512:
-  metadata.gz: 2a9ab36b64b600bbd3f306d818f2bfc06666e520bba5588614ed5783c78227b1277522f03737de90b185ef659bd099cbb4064ece4d2afbafe5e8815e71916ca8
-  data.tar.gz: a94a6ca7fddac35799d4dad64099aacc0af68b4d7e37ffcd777d504852a29f26016595d3e5a15c3a3a9174a9706709787342f80fb059e4a01292d00f219e8df6
+  metadata.gz: a97e4828c3c0d41607c021fe3b3b87b9c2c2dc01ae70131a984ce4513a46c2386e8d86ce5fc3029f395cc16a9f410c1fbe4c2dff8f77c833bc40f4d75ff8e509
+  data.tar.gz: def00e2668b59de09ed0eb8a2fa8efb5974c25c36ea82fcf64b83115afb7519b38010019310267dc4efddeca76df56e6861696abd285d2b2dfab8c00a309e21b

data/lib/pg/replication/buffer.rb

@@ -1,21 +1,49 @@
 # frozen_string_literal: true
 
-require "delegate"
-require "stringio"
-
 module PG
   module Replication
-    class Buffer < SimpleDelegator
+    POSTGRES_EPOCH = Time.utc(2000, 1, 1).freeze
+    POSTGRES_EPOCH_USECS = (POSTGRES_EPOCH.to_r * 1_000_000).to_i
+
+    class Buffer
+      def initialize(data)
+        @data = data
+        @pos = 0
+      end
+
       def self.from_string(str)
-        new(StringIO.new(str))
+        new(str.b)
+      end
+
+      def eof?
+        @pos >= @data.bytesize
+      end
+
+      def read(n = nil)
+        return nil if @pos >= @data.bytesize
+        if n.nil?
+          result = @data.byteslice(@pos..-1)
+          @pos = @data.bytesize
+        else
+          result = @data.byteslice(@pos, n)
+          @pos += result.bytesize
+        end
+        result
+      end
+
+      def readbyte
+        raise EOFError if @pos >= @data.bytesize
+        byte = @data.getbyte(@pos)
+        @pos += 1
+        byte
       end
 
       def read_char
-        read_int8.chr
+        readbyte.chr
       end
 
       def read_bool
-        read_int8 == 1
+        readbyte == 1
       end
 
       def read_int8
@@ -23,40 +51,38 @@ module PG
       end
 
       def read_int16
-        read_bytes(2).unpack("n").first
+        read_bytes(2).unpack1("n")
       end
 
       def read_int32
-        read_bytes(4).unpack("N").first
+        read_bytes(4).unpack1("N")
       end
 
       def read_int64
-        read_bytes(8).unpack("Q>").first
+        read_bytes(8).unpack1("Q>")
       end
 
       def read_timestamp
-        usecs = Time.new(2_000, 1, 1, 0, 0, 0, 0).to_i * 10**6 + read_int64
-        Time.at(usecs / 10**6, usecs % 10**6, :microsecond)
+        usecs = POSTGRES_EPOCH_USECS + read_int64
+        Time.at(usecs / 1_000_000, usecs % 1_000_000, :microsecond, in: "UTC")
       end
 
       def read_cstring
-        str = String.new
-        loop do
-          case read_char
-          in "\0"
-            return str
-          in chr
-            str << chr
-          end
-        end
+        null_pos = @data.index("\0", @pos)
+        raise EOFError, "Unterminated C-string" if null_pos.nil?
+
+        str = @data.byteslice(@pos, null_pos - @pos)
+        @pos = null_pos + 1 # Skip past null terminator
+        str
       end
 
       private
 
       def read_bytes(n)
-        bytes = read(n)
-        raise EOFError if bytes.nil? || bytes.size < n
-        bytes
+        raise EOFError if @pos + n > @data.bytesize
+        result = @data.byteslice(@pos, n)
+        @pos += n
+        result
       end
     end
   end

data/lib/pg/replication/version.rb

@@ -2,6 +2,6 @@
 
 module PG
   module Replication
-    VERSION = "0.0.7"
+    VERSION = "0.0.8"
   end
 end

data/lib/pg/replication.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "pg"
+require "thread"
 require_relative "replication/version"
 require_relative "replication/buffer"
 require_relative "replication/pg_output"
@@ -8,7 +9,12 @@ require_relative "replication/protocol"
 
 module PG
   module Replication
-    def start_replication_slot(slot, logical: true, auto_keep_alive: true, location: "0/0", **params)
+    DEFAULT_QUEUE_SIZE = 10_000
+
+    StreamEnd = Object.new.freeze
+    StreamError = Data.define(:exception)
+
+    def start_replication_slot(slot, logical: true, auto_keep_alive: true, location: "0/0", queue_size: DEFAULT_QUEUE_SIZE, **params)
       keep_alive_secs = wal_receiver_status_interval
       @last_confirmed_lsn = confirmed_slot_lsn(slot) || 0
 
@@ -22,51 +28,11 @@ module PG
       end
       query(start_query)
 
-      last_keep_alive = Time.now
-
-      Enumerator
-        .new do |y|
-          loop do
-            if auto_keep_alive && Time.now - last_keep_alive > keep_alive_secs
-              standby_status_update(write_lsn: @last_confirmed_lsn)
-              last_keep_alive = Time.now
-            end
-
-            consume_input
-            next if is_busy
-
-            case get_copy_data(async: true)
-            in nil
-              get_last_result
-              break
-
-            in false
-              IO.select([socket_io], nil, nil, keep_alive_secs)
-              next
-
-            in data
-              case (msg = Protocol.read_message(Buffer.new(StringIO.new(data))))
-              in Protocol::XLogData(lsn:, data:) if auto_keep_alive
-                y << msg
-                standby_status_update(write_lsn: lsn) if lsn > 0
-                last_keep_alive = Time.now
-
-              in Protocol::PrimaryKeepalive(current_lsn:, server_time:, asap: true) if auto_keep_alive
-                standby_status_update(write_lsn: current_lsn)
-                last_keep_alive = Time.now
-                y << msg
-
-              in Protocol::PrimaryKeepalive(current_lsn:)
-                y << msg
-                @last_confirmed_lsn = [@last_confirmed_lsn, current_lsn].compact.max
-
-              else
-                y << msg
-              end
-            end
-          end
-        end
-        .lazy
+      if auto_keep_alive
+        start_threaded_replication(keep_alive_secs, queue_size)
+      else
+        start_sync_replication
+      end
     end
 
     def start_pgoutput_replication_slot(slot, publication_names, **kwargs)
@@ -96,17 +62,25 @@ module PG
         write_lsn,
         flush_lsn,
         apply_lsn,
-        (timestamp - Time.new(2_000, 1, 1, 0, 0, 0, 0)) * 10**6,
+        ((timestamp.to_r - POSTGRES_EPOCH.to_r) * 1_000_000).to_i,
         reply ? 1 : 0,
       ].pack("CQ>Q>Q>Q>C")
 
-      put_copy_data(msg)
-      flush
-      @last_confirmed_lsn = [@last_confirmed_lsn, write_lsn].compact.max
+      status_update_mutex.synchronize do
+        put_copy_data(msg)
+        flush
+        @last_confirmed_lsn = [@last_confirmed_lsn, write_lsn].compact.max
+      end
     end
 
     def last_confirmed_lsn
-      @last_confirmed_lsn
+      status_update_mutex.synchronize { @last_confirmed_lsn }
+    end
+
+    def stop_replication
+      status_update_mutex.synchronize do
+        put_copy_end
+      end
     end
 
     def wal_receiver_status_interval
@@ -124,6 +98,120 @@ module PG
     rescue StandardError
       nil
     end
+
+    private
+
+    def status_update_mutex
+      @status_update_mutex ||= Mutex.new
+    end
+
+    def start_threaded_replication(keep_alive_secs, queue_size)
+      conn = self
+      queue = SizedQueue.new(queue_size)
+      last_keepalive = Time.now
+
+      thread = Thread.new do
+        loop do
+          break if queue.closed?
+
+          conn.consume_input
+          next if conn.is_busy
+
+          case (data = conn.get_copy_data(async: true))
+          when nil
+            queue.push(StreamEnd, true) rescue nil
+            conn.get_last_result
+            break
+
+          when false
+            timeout = [keep_alive_secs - (Time.now - last_keepalive), 0.1].max
+            IO.select([conn.socket_io], nil, nil, timeout)
+
+            if Time.now - last_keepalive >= keep_alive_secs
+              conn.standby_status_update(write_lsn: conn.last_confirmed_lsn)
+              last_keepalive = Time.now
+            end
+
+          else
+            msg = Protocol.read_message(Buffer.from_string(data))
+
+            if msg.is_a?(Protocol::PrimaryKeepalive) && msg.asap
+              conn.standby_status_update(write_lsn: msg.current_lsn)
+              last_keepalive = Time.now
+            end
+
+            # Non-blocking push with keepalive retry loop
+            loop do
+              break if queue.closed?
+              begin
+                queue.push(msg, true)
+                break
+              rescue ThreadError
+                if Time.now - last_keepalive >= keep_alive_secs
+                  conn.standby_status_update(write_lsn: conn.last_confirmed_lsn)
+                  last_keepalive = Time.now
+                end
+                sleep(0.05)
+              end
+            end
+          end
+        end
+      rescue ClosedQueueError
+        # Clean exit
+      rescue => e
+        queue.push(StreamError.new(e), true) rescue nil
+      end
+
+      Enumerator.new do |y|
+        loop do
+          msg = queue.pop
+
+          case msg
+          when StreamEnd
+            break
+          when StreamError
+            raise msg.exception
+          else
+            y << msg
+
+            lsn = case msg
+            when Protocol::XLogData
+              msg.lsn
+            when Protocol::PrimaryKeepalive
+              msg.current_lsn
+            end
+
+            if lsn && lsn > 0
+              status_update_mutex.synchronize { @last_confirmed_lsn = lsn }
+            end
+          end
+        end
+      ensure
+        queue.close
+        thread.join(5)
+        thread.kill if thread.alive?
+      end.lazy
+    end
+
+    def start_sync_replication
+      Enumerator.new do |y|
+        loop do
+          consume_input
+          next if is_busy
+
+          case get_copy_data(async: true)
+          in nil
+            get_last_result
+            break
+          in false
+            IO.select([socket_io], nil, nil, 10)
+            next
+          in data
+            y << Protocol.read_message(Buffer.new(data))
+          end
+        end
+      end.lazy
+    end
   end
 
   Connection.send(:include, Replication)

metadata

@@ -1,13 +1,13 @@
 --- !ruby/object:Gem::Specification
 name: pg-replication-protocol
 version: !ruby/object:Gem::Version
-  version: 0.0.7
+  version: 0.0.8
 platform: ruby
 authors:
 - Rodrigo Navarro
 bindir: bin
 cert_chain: []
-date: 2025-06-22 00:00:00.000000000 Z
+date: 2026-01-14 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: pg