From 2448703414d9f516cc7e06f7446e76600545d949 Mon Sep 17 00:00:00 2001
From: Nathan Zimmerberg <39104088+nhz2@users.noreply.github.com>
Date: Tue, 17 Dec 2024 21:29:49 -0500
Subject: [PATCH] Add finalizer to BZStream and make
 `TranscodingStreams.finalize` a noop (#43)

* add finalizer to BZStream

* add basic test for memory leak
---
 src/compression.jl    | 21 +++++----------------
 src/decompression.jl  | 21 +++++----------------
 src/libbzip2.jl       | 10 ++++++++++
 test/big-mem-tests.jl | 10 ++++++++++
 test/runtests.jl      |  7 +++++++
 5 files changed, 37 insertions(+), 32 deletions(-)

diff --git a/src/compression.jl b/src/compression.jl
index b6da111..b21391a 100644
--- a/src/compression.jl
+++ b/src/compression.jl
@@ -1,7 +1,7 @@
 # Compressor Codec
 # ================
 
-mutable struct Bzip2Compressor <: TranscodingStreams.Codec
+struct Bzip2Compressor <: TranscodingStreams.Codec
     stream::BZStream
     blocksize100k::Int
     workfactor::Int
@@ -37,7 +37,9 @@ function Bzip2Compressor(;blocksize100k::Integer=DEFAULT_BLOCKSIZE100K,
     elseif !(0 ≤ verbosity ≤ 4)
         throw(ArgumentError("verbosity must be within 0..4"))
     end
-    return Bzip2Compressor(BZStream(), blocksize100k, workfactor, verbosity)
+    stream = BZStream()
+    finalizer(compress_finalizer!, stream)
+    return Bzip2Compressor(stream, blocksize100k, workfactor, verbosity)
 end
 
 const Bzip2CompressorStream{S} = TranscodingStream{Bzip2Compressor,S} where S<:IO
@@ -56,23 +58,10 @@ end
 # Methods
 # -------
 
-function TranscodingStreams.finalize(codec::Bzip2Compressor)
-    if codec.stream.state != C_NULL
-        code = compress_end!(codec.stream)
-        if code != BZ_OK
-            bzerror(code)
-        end
-    end
-    return
-end
-
 function TranscodingStreams.startproc(codec::Bzip2Compressor, ::Symbol, error_ref::Error)
     if codec.stream.state != C_NULL
         code = compress_end!(codec.stream)
-        if code != BZ_OK
-            error_ref[] = BZ2Error(code)
-            return :error
-        end
+        @assert code == BZ_OK
     end
     code = compress_init!(codec.stream, codec.blocksize100k, codec.verbosity, codec.workfactor)
     # errors in compress_init! do not require clean up, so just throw
diff --git a/src/decompression.jl b/src/decompression.jl
index 8582b73..40ecf24 100644
--- a/src/decompression.jl
+++ b/src/decompression.jl
@@ -1,7 +1,7 @@
 # Decompressor Codec
 # ==================
 
-mutable struct Bzip2Decompressor <: TranscodingStreams.Codec
+struct Bzip2Decompressor <: TranscodingStreams.Codec
     stream::BZStream
     small::Bool
     verbosity::Int
@@ -25,7 +25,9 @@ function Bzip2Decompressor(;small::Bool=false, verbosity::Integer=DEFAULT_VERBOS
     if !(0 ≤ verbosity ≤ 4)
         throw(ArgumentError("verbosity must be within 0..4"))
     end
-    return Bzip2Decompressor(BZStream(), small, verbosity)
+    stream = BZStream()
+    finalizer(decompress_finalizer!, stream)
+    return Bzip2Decompressor(stream, small, verbosity)
 end
 
 const Bzip2DecompressorStream{S} = TranscodingStream{Bzip2Decompressor,S} where S<:IO
@@ -44,23 +46,10 @@ end
 # Methods
 # -------
 
-function TranscodingStreams.finalize(codec::Bzip2Decompressor)
-    if codec.stream.state != C_NULL
-        code = decompress_end!(codec.stream)
-        if code != BZ_OK
-            bzerror(code)
-        end
-    end
-    return
-end
-
 function TranscodingStreams.startproc(codec::Bzip2Decompressor, ::Symbol, error_ref::Error)
     if codec.stream.state != C_NULL
         code = decompress_end!(codec.stream)
-        if code != BZ_OK
-            error_ref[] = BZ2Error(code)
-            return :error
-        end
+        @assert code == BZ_OK
     end
     code = decompress_init!(codec.stream, codec.verbosity, codec.small)
     # errors in decompress_init! do not require clean up, so just throw
diff --git a/src/libbzip2.jl b/src/libbzip2.jl
index 48343f2..951564b 100644
--- a/src/libbzip2.jl
+++ b/src/libbzip2.jl
@@ -106,6 +106,11 @@ function compress_end!(stream::BZStream)
     end
 end
 
+function compress_finalizer!(stream::BZStream)
+    compress_end!(stream)
+    nothing
+end
+
 function compress!(stream::BZStream, action::Integer)
     if WIN32
         return ccall(
@@ -161,6 +166,11 @@ function decompress_end!(stream::BZStream)
     end
 end
 
+function decompress_finalizer!(stream::BZStream)
+    decompress_end!(stream)
+    nothing
+end
+
 function decompress!(stream::BZStream)
     if WIN32
         return ccall(
diff --git a/test/big-mem-tests.jl b/test/big-mem-tests.jl
index 0856182..d120c27 100644
--- a/test/big-mem-tests.jl
+++ b/test/big-mem-tests.jl
@@ -6,6 +6,16 @@
 using Test
 using CodecBzip2
 
+@testset "memory leak" begin
+    function foo()
+        for i in 1:1000000
+            c = transcode(Bzip2Compressor(), zeros(UInt8,16))
+            u = transcode(Bzip2Decompressor(), c)
+        end
+    end
+    foo()
+end
+
 @testset "Big Memory Tests" begin
     Sys.WORD_SIZE == 64 || error("tests require 64 bit word size")
     @info "compressing zeros"
diff --git a/test/runtests.jl b/test/runtests.jl
index e49198e..5f1a912 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -86,4 +86,11 @@ Aqua.test_all(CodecBzip2)
         @test sprint(Base.showerror, CodecBzip2.BZ2Error(-100)) ==
             "BZ2Error: unknown bzip2 error code: -100"
     end
+    @testset "memory leaks" begin
+        # issue #27
+        for i in 1:200000
+            c = transcode(Bzip2Compressor(), zeros(UInt8,16))
+            u = transcode(Bzip2Decompressor(), c)
+        end
+    end
 end