pg-replication-protocol 0.0.6 → 0.0.8

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 72bf40ef6c6c482636553033a8747eec8239f4b10602486376806fd6bff59a91
-  data.tar.gz: 4df39c13ee3b16f7e77114c7ef061e75b33b1f112e6341ce4aeb28898096d703
+  metadata.gz: f00c769aee5ae5e1d41063e827f2d5fcaa19991aed6d4676f414ce5f6d6326ef
+  data.tar.gz: a99c6f062f9c33109a409b58ef465833fb1d0b3fd8499a91461db9ea3e2586af
 SHA512:
-  metadata.gz: 6b3a4cd7627a4d8229f88ba7d086a95a6fac1ae76fe7d856970ce1dd745272696046f18e55bb551345fa8b431691ac351d3cffcdc809f07ecba5fa78d6f4b67a
-  data.tar.gz: da57dba3124e0c3e8c5832d570f471b2a15ee68c2afc0fd71e0f5c5410ec61585f99dcfb1499059106cf1cb06073552335897035388b45aa884dfc0d757155e0
+  metadata.gz: a97e4828c3c0d41607c021fe3b3b87b9c2c2dc01ae70131a984ce4513a46c2386e8d86ce5fc3029f395cc16a9f410c1fbe4c2dff8f77c833bc40f4d75ff8e509
+  data.tar.gz: def00e2668b59de09ed0eb8a2fa8efb5974c25c36ea82fcf64b83115afb7519b38010019310267dc4efddeca76df56e6861696abd285d2b2dfab8c00a309e21b

data/Gemfile

@@ -5,6 +5,8 @@ source "https://rubygems.org"
 gemspec
 
 gem "rake", "~> 13.0"
+gem "base64"
+gem "rbs"
 
 group :test do
   gem "rspec", "~> 3.0"

data/Gemfile.lock

@@ -1,20 +1,24 @@
 PATH
   remote: .
   specs:
-    pg-replication-protocol (0.0.5)
+    pg-replication-protocol (0.0.7)
       pg (~> 1.0)
 
 GEM
   remote: https://rubygems.org/
   specs:
+    base64 (0.3.0)
     diff-lcs (1.5.1)
     docker-api (2.4.0)
       excon (>= 0.64.0)
       multi_json
     excon (1.2.1)
+    logger (1.7.0)
     multi_json (1.15.0)
     pg (1.5.9)
     rake (13.2.1)
+    rbs (3.9.4)
+      logger
     rspec (3.13.0)
       rspec-core (~> 3.13.0)
       rspec-expectations (~> 3.13.0)
@@ -38,8 +42,10 @@ PLATFORMS
   x86_64-linux
 
 DEPENDENCIES
+  base64
   pg-replication-protocol!
   rake (~> 13.0)
+  rbs
   rspec (~> 3.0)
   testcontainers-postgres
 

data/Rakefile

@@ -2,3 +2,7 @@
 
 require "bundler/gem_tasks"
 task default: %i[]
+
+task :test do
+  sh "bundle exec rbs test --target PG::Replication::* bundle exec rspec"
+end

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.6"
+    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)

data/sig/pg/replication.rbs

@@ -0,0 +1,283 @@
+module PG
+  module Replication
+    module Protocol
+      class PrimaryKeepalive < ::Data
+        def self.new: (Integer current_lsn, Time server_time, Boolean asap) -> instance
+                    | (current_lsn: Integer, server_time: Time, asap: Boolean) -> instance
+
+        def self.[]: (Integer current_lsn, Time server_time, Boolean asap) -> instance
+                   | (current_lsn: Integer, server_time: Time, asap: Boolean) -> instance
+
+        def self.members: () -> [ :current_lsn, :server_time, :asap ]
+
+        def members: () -> [ :current_lsn, :server_time, :asap ]
+
+        attr_reader current_lsn: Integer
+
+        attr_reader server_time: Time
+
+        attr_reader asap: Boolean
+      end
+
+      class XLogData[T] < ::Data
+        def self.new: (Integer lsn, Integer current_lsn, Time server_time, T data) -> instance
+                    | (lsn: Integer, current_lsn: Integer, server_time: Time, data: T) -> instance
+
+        def self.[]: (Integer lsn, Integer current_lsn, Time server_time, T data) -> instance
+                   | (lsn: Integer, current_lsn: Integer, server_time: Time, data: T) -> instance
+
+        def self.members: () -> [ :lsn, :current_lsn, :server_time, :data ]
+
+        def members: () -> [ :lsn, :current_lsn, :server_time, :data ]
+
+        attr_reader lsn: Integer
+
+        attr_reader current_lsn: Integer
+
+        attr_reader server_time: Time
+
+        attr_reader data: T
+      end
+    end
+
+    module PGOutput
+      class Begin < ::Data
+        def self.new: (Integer final_lsn, Time timestamp, Integer xid) -> instance
+                    | (final_lsn: Integer, timestamp: Time, xid: Integer) -> instance
+
+        def self.[]: (Integer final_lsn, Time timestamp, Integer xid) -> instance
+                   | (final_lsn: Integer, timestamp: Time, xid: Integer) -> instance
+
+        def self.members: () -> [ :final_lsn, :timestamp, :xid ]
+
+        def members: () -> [ :final_lsn, :timestamp, :xid ]
+
+        attr_reader final_lsn: Integer
+
+        attr_reader timestamp: Time
+
+        attr_reader xid: Integer
+      end
+
+      class Column < ::Data
+        def self.new: (Integer flags, String name, Integer oid, Integer modifier) -> instance
+                    | (flags: Integer, name: String, oid: Integer, modifier: Integer) -> instance
+
+        def self.[]: (Integer flags, String name, Integer oid, Integer modifier) -> instance
+                   | (flags: Integer, name: String, oid: Integer, modifier: Integer) -> instance
+
+        def self.members: () -> [ :flags, :name, :oid, :modifier ]
+
+        def members: () -> [ :flags, :name, :oid, :modifier ]
+
+        attr_reader flags: Integer
+
+        attr_reader name: String
+
+        attr_reader oid: Integer
+
+        attr_reader modifier: Integer
+
+        def key?: () -> Boolean
+      end
+
+      class Commit < ::Data
+        def self.new: (Integer lsn, Integer end_lsn, Time timestamp) -> instance
+                    | (lsn: Integer, end_lsn: Integer, timestamp: Time) -> instance
+
+        def self.[]: (Integer lsn, Integer end_lsn, Time timestamp) -> instance
+                   | (lsn: Integer, end_lsn: Integer, timestamp: Time) -> instance
+
+        def self.members: () -> [ :lsn, :end_lsn, :timestamp ]
+
+        def members: () -> [ :lsn, :end_lsn, :timestamp ]
+
+        attr_reader lsn: Integer
+
+        attr_reader end_lsn: Integer
+
+        attr_reader timestamp: Time
+      end
+
+      class Delete < ::Data
+        def self.new: (Integer oid, Array[String] key, Array[String] old) -> instance
+                    | (oid: Integer, key: Array[String], old: Array[String]) -> instance
+
+        def self.[]: (Integer oid, Array[String] key, Array[String] old) -> instance
+                   | (oid: Integer, key: Array[String], old: Array[String]) -> instance
+
+        def self.members: () -> [ :oid, :key, :old ]
+
+        def members: () -> [ :oid, :key, :old ]
+
+        attr_reader oid: Integer
+
+        attr_reader key: Array[String]
+
+        attr_reader old: Array[String]
+      end
+
+      class Insert < ::Data
+        def self.new: (Integer oid, Array[String] new) -> instance
+                    | (oid: Integer, new: Array[String]) -> instance
+
+        def self.[]: (Integer oid, Array[String] new) -> instance
+                   | (oid: Integer, new: Array[String]) -> instance
+
+        def self.members: () -> [ :oid, :new ]
+
+        def members: () -> [ :oid, :new ]
+
+        attr_reader oid: Integer
+
+        attr_reader new: Array[String]
+      end
+
+      class Message < ::Data
+        def self.new: (Boolean transactional, Integer lsn, String? prefix, String? content) -> instance
+                    | (transactional: Boolean, lsn: Integer, prefix: String?, content: String?) -> instance
+
+        def self.[]: (Boolean transactional, Integer lsn, String? prefix, String? content) -> instance
+                   | (transactional: Boolean, lsn: Integer, prefix: String?, content: String?) -> instance
+
+        def self.members: () -> [ :transactional, :lsn, :prefix, :content ]
+
+        def members: () -> [ :transactional, :lsn, :prefix, :content ]
+
+        attr_reader transactional: Boolean
+
+        attr_reader lsn: Integer
+
+        attr_reader prefix: String?
+
+        attr_reader content: String?
+      end
+
+      class Origin < ::Data
+        def self.new: (Integer commit_lsn, String name) -> instance
+                    | (commit_lsn: Integer, name: String) -> instance
+
+        def self.[]: (Integer commit_lsn, String name) -> instance
+                   | (commit_lsn: Integer, name: String) -> instance
+
+        def self.members: () -> [ :commit_lsn, :name ]
+
+        def members: () -> [ :commit_lsn, :name ]
+
+        attr_reader commit_lsn: Integer
+
+        attr_reader name: String
+      end
+
+      class Relation < ::Data
+        def self.new: (Integer oid, String namespace, String name, String replica_identity, Array[PG::Replication::PGOutput::Column] columns) -> instance
+                    | (oid: Integer, namespace: String, name: String, replica_identity: String, columns: Array[PG::Replication::PGOutput::Column]) -> instance
+
+        def self.[]: (Integer oid, String namespace, String name, String replica_identity, Array[PG::Replication::PGOutput::Column] columns) -> instance
+                   | (oid: Integer, namespace: String, name: String, replica_identity: String, columns: Array[PG::Replication::PGOutput::Column]) -> instance
+
+        def self.members: () -> [ :oid, :namespace, :name, :replica_identity, :columns ]
+
+        def members: () -> [ :oid, :namespace, :name, :replica_identity, :columns ]
+
+        attr_reader oid: Integer
+
+        attr_reader namespace: String
+
+        attr_reader name: String
+
+        attr_reader replica_identity: String
+
+        attr_reader columns: Array[PG::Replication::PGOutput::Column]
+      end
+
+      class Truncate < ::Data
+        def self.new: (Integer oid) -> instance
+                    | (oid: Integer) -> instance
+
+        def self.[]: (Integer oid) -> instance
+                   | (oid: Integer) -> instance
+
+        def self.members: () -> [ :oid ]
+
+        def members: () -> [ :oid ]
+
+        attr_reader oid: Integer
+      end
+
+      class Tuple < ::Data
+        def self.new: (String type, String? data) -> instance
+                    | (type: String, data: String?) -> instance
+
+        def self.[]: (String type, String? data) -> instance
+                   | (type: String, data: String?) -> instance
+
+        def self.members: () -> [ :type, :data ]
+
+        def members: () -> [ :type, :data ]
+
+        attr_reader type: String
+
+        attr_reader data: String?
+
+        def text?: () -> Boolean
+
+        def binary?: () -> Boolean
+
+        def toast?: () -> Boolean
+      end
+
+      class Type < ::Data
+        def self.new: (Integer oid, String namespace, String name) -> instance
+                    | (oid: Integer, namespace: String, name: String) -> instance
+
+        def self.[]: (Integer oid, String namespace, String name) -> instance
+                   | (oid: Integer, namespace: String, name: String) -> instance
+
+        def self.members: () -> [ :oid, :namespace, :name ]
+
+        def members: () -> [ :oid, :namespace, :name ]
+
+        attr_reader oid: Integer
+
+        attr_reader namespace: String
+
+        attr_reader name: String
+      end
+
+      class Update < ::Data
+        def self.new: (Integer oid, Array[String] key, Array[String] old, Array[String] new) -> instance
+                    | (oid: Integer, key: Array[String], old: Array[String], new: Array[String]) -> instance
+
+        def self.[]: (Integer oid, Array[String] key, Array[String] old, Array[String] new) -> instance
+                   | (oid: Integer, key: Array[String], old: Array[String], new: Array[String]) -> instance
+
+        def self.members: () -> [ :oid, :key, :old, :new ]
+
+        def members: () -> [ :oid, :key, :old, :new ]
+
+        attr_reader oid: Integer
+
+        attr_reader key: Array[String]
+
+        attr_reader old: Array[String]
+
+        attr_reader new: Array[String]
+      end
+    end
+
+    @last_confirmed_lsn: Integer?
+
+    def start_replication_slot: (String slot, ?logical: bool, ?auto_keep_alive: bool, ?location: ::String, **untyped params) -> Enumerator::Lazy[PG::Replication::Protocol::PrimaryKeepalive | PG::Replication::Protocol::XLogData[String]]
+
+    def start_pgoutput_replication_slot: (String slot, Array[String] publication_names, **untyped kwargs) -> Enumerator::Lazy[PG::Replication::Protocol::PrimaryKeepalive | PG::Replication::Protocol::XLogData[PG::Replication::PGOutput]]
+
+    def standby_status_update: (write_lsn: Integer, ?flush_lsn: Integer, ?apply_lsn: Integer, ?timestamp: Time, ?reply: bool) -> void
+
+    def last_confirmed_lsn: () -> Integer?
+
+    def wal_receiver_status_interval: () -> Integer
+
+    def confirmed_slot_lsn: (String slot) -> Integer?
+  end
+end

metadata

@@ -1,13 +1,13 @@
 --- !ruby/object:Gem::Specification
 name: pg-replication-protocol
 version: !ruby/object:Gem::Version
-  version: 0.0.6
+  version: 0.0.8
 platform: ruby
 authors:
 - Rodrigo Navarro
 bindir: bin
 cert_chain: []
-date: 2025-06-20 00:00:00.000000000 Z
+date: 2026-01-14 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: pg
@@ -41,6 +41,7 @@ files:
 - lib/pg/replication/protocol.rb
 - lib/pg/replication/version.rb
 - pg-replication-protocol.gemspec
+- sig/pg/replication.rbs
 licenses:
 - MIT
 metadata: