Refactor memory management in BAM and BCF classes to ensure proper buffer handling and prevent memory leaks

Author: kojix2

lib/hts/bam/header.rb

@@ -81,15 +81,23 @@ def <<(obj)
       # experimental
       def find_line(type, key, value)
         ks = LibHTS::KString.new
-        r = LibHTS.sam_hdr_find_line_id(@sam_hdr, type, key, value, ks)
-        r == 0 ? ks[:s] : nil
+        begin
+          r = LibHTS.sam_hdr_find_line_id(@sam_hdr, type, key, value, ks)
+          r == 0 ? ks.read_string_copy : nil
+        ensure
+          ks.free_buffer
+        end
       end
 
       # experimental
       def find_line_at(type, pos)
         ks = LibHTS::KString.new
-        r = LibHTS.sam_hdr_find_line_pos(@sam_hdr, type, pos, ks)
-        r == 0 ? ks[:s] : nil
+        begin
+          r = LibHTS.sam_hdr_find_line_pos(@sam_hdr, type, pos, ks)
+          r == 0 ? ks.read_string_copy : nil
+        ensure
+          ks.free_buffer
+        end
       end
 
       # experimental

lib/hts/bam/record.rb

@@ -354,9 +354,13 @@ def base_mod(auto_parse: true)
       # @return [String] a string representation of the alignment.
       def to_s
         kstr = LibHTS::KString.new
-        raise "Failed to format bam record" if LibHTS.sam_format1(@header.struct, @bam1, kstr) == -1
+        begin
+          raise "Failed to format bam record" if LibHTS.sam_format1(@header.struct, @bam1, kstr) == -1
 
-        kstr[:s]
+          kstr.read_string_copy
+        ensure
+          kstr.free_buffer
+        end
       end
 
       private

lib/hts/bcf/header.rb

@@ -85,15 +85,23 @@ def get_hrec(bcf_hl_type, key, value, str_class = nil)
       def seqnames
         n = FFI::MemoryPointer.new(:int)
         names = LibHTS.bcf_hdr_seqnames(@bcf_hdr, n)
-        names.read_array_of_pointer(n.read_int)
-             .map(&:read_string)
+        begin
+          names.read_array_of_pointer(n.read_int)
+               .map(&:read_string)
+        ensure
+          LibHTS.hts_free(names) unless names.null?
+        end
       end
 
       def to_s
         kstr = LibHTS::KString.new
-        raise "Failed to get header string" unless LibHTS.bcf_hdr_format(@bcf_hdr, 0, kstr)
+        begin
+          raise "Failed to get header string" unless LibHTS.bcf_hdr_format(@bcf_hdr, 0, kstr)
 
-        kstr[:s]
+          kstr.read_string_copy
+        ensure
+          kstr.free_buffer
+        end
       end
 
       def name2id(name)

lib/hts/bcf/header_record.rb

@@ -31,8 +31,12 @@ def find_key(key)
 
       def to_s
         kstr = LibHTS::KString.new
-        LibHTS.bcf_hrec_format(@bcf_hrec, kstr)
-        kstr[:s]
+        begin
+          LibHTS.bcf_hrec_format(@bcf_hrec, kstr)
+          kstr.read_string_copy
+        ensure
+          kstr.free_buffer
+        end
       end
 
       private

lib/hts/bcf/record.rb

@@ -130,9 +130,13 @@ def format(key = nil)
 
       def to_s
         ksr = LibHTS::KString.new
-        raise "Failed to format record" if LibHTS.vcf_format(@header.struct, @bcf1, ksr) == -1
+        begin
+          raise "Failed to format record" if LibHTS.vcf_format(@header.struct, @bcf1, ksr) == -1
 
-        ksr[:s]
+          ksr.read_string_copy
+        ensure
+          ksr.free_buffer
+        end
       end
 
       private

lib/hts/libhts/constants.rb

@@ -12,6 +12,27 @@ class KString < FFI::Struct
         :l,              :size_t,
         :m,              :size_t,
         :s,              :string
+
+      def buffer_ptr
+        to_ptr.get_pointer(self.class.offset_of(:s))
+      end
+
+      def read_string_copy
+        ptr = buffer_ptr
+        return "" if ptr.null?
+
+        ptr.read_string(self[:l])
+      end
+
+      def free_buffer
+        ptr = buffer_ptr
+        return if ptr.null?
+
+        LibHTS.hts_free(ptr)
+        to_ptr.put_pointer(self.class.offset_of(:s), FFI::Pointer::NULL)
+        self[:l] = 0
+        self[:m] = 0
+      end
     end
 
     class KSeq < FFI::Struct

lib/hts/tabix.rb

@@ -89,8 +89,11 @@ def name2id(name)
     def seqnames
       check_closed
       nseq = FFI::MemoryPointer.new(:int)
-      LibHTS.tbx_seqnames(@idx, nseq).then do |pts|
+      pts = LibHTS.tbx_seqnames(@idx, nseq)
+      begin
         pts.read_array_of_pointer(nseq.read_int).map(&:read_string)
+      ensure
+        LibHTS.hts_free(pts) unless pts.null?
       end
     end
 
@@ -144,8 +147,11 @@ def querys(region, &block)
     def query_yield(qiter)
       r = LibHTS::KString.new
       begin
-        yield r[:s].split("\t") while LibHTS.tbx_itr_next(@hts_file, @idx, qiter, r) > 0
+        while LibHTS.tbx_itr_next(@hts_file, @idx, qiter, r) > 0
+          yield r.read_string_copy.split("\t")
+        end
       ensure
+        r.free_buffer
         LibHTS.hts_itr_destroy(qiter)
       end
     end