Only reset session on error in process, override transcode with pledge

src/compression.jl

@@ -77,6 +77,25 @@ end
 
 # Methods
 # -------
+function TranscodingStreams.transcode(
+    codec::ZstdCompressor,
+    input::TranscodingStreams.Buffer,
+    output::Union{TranscodingStreams.Buffer,Nothing} = nothing,
+)
+    code = LibZstd.ZSTD_CCtx_setPledgedSrcSize(codec.cstream, TranscodingStreams.buffersize(input))
+    if iserror(code)
+        zstderror(codec.cstream, code)
+    end
+    return invoke(
+        TranscodingStreams.transcode,
+        Tuple{
+              TranscodingStreams.Codec,
+              TranscodingStreams.Buffer,
+              typeof(output)
+        },
+        codec, input, output
+    )
+end
 
 function TranscodingStreams.initialize(codec::ZstdCompressor)
     code = initialize!(codec.cstream, codec.level)
@@ -102,11 +121,8 @@ function TranscodingStreams.finalize(codec::ZstdCompressor)
 end
 
 function TranscodingStreams.startproc(codec::ZstdCompressor, mode::Symbol, error::Error)
-    code = reset!(codec.cstream, 0 #=unknown source size=#)
-    if iserror(code)
-        error[] = ErrorException("zstd error")
-        return :error
-    end
+    reset!(codec.cstream.ibuffer)
+    reset!(codec.cstream.obuffer)
     return :ok
 end
 
@@ -146,6 +162,7 @@ function TranscodingStreams.process(codec::ZstdCompressor, input::Memory, output
     if iserror(code)
         ptr = LibZstd.ZSTD_getErrorName(code)
         error[] = ErrorException("zstd error: " * unsafe_string(ptr))
+        reset!(cstream)
         return Δin, Δout, :error
     else
         return Δin, Δout, input.size == 0 && code == 0 ? :end : :ok

src/libzstd.jl

@@ -60,6 +60,12 @@ function initialize!(cstream::CStream, level::Integer)
     return LibZstd.ZSTD_initCStream(cstream, level)
 end
 
+function reset!(cstream::CStream)
+    # Resetting session never fails.
+    # Also this will reset the pledged size
+    return LibZstd.ZSTD_CCtx_reset(cstream, LibZstd.ZSTD_reset_session_only)
+end
+
 function reset!(cstream::CStream, srcsize::Integer)
     # ZSTD_resetCStream is deprecated
     # https://github.com/facebook/zstd/blob/9d2a45a705e22ad4817b41442949cd0f78597154/lib/zstd.h#L2253-L2272
@@ -163,6 +169,11 @@ end
 const ZSTD_CONTENTSIZE_UNKNOWN = Culonglong(0) - 1
 const ZSTD_CONTENTSIZE_ERROR   = Culonglong(0) - 2
 
+# ZSTD_findDecompressedSize gets the decompressed size of all available frames
+# Alternatively, consider ZSTD_getFrameContentSize for a single frame
 function find_decompressed_size(src::Ptr, size::Integer)
     return LibZstd.ZSTD_findDecompressedSize(src, size)
 end
+function find_decompressed_size(src::Vector{UInt8})
+    return LibZstd.ZSTD_findDecompressedSize(src, length(src))
+end

test/runtests.jl

@@ -38,6 +38,7 @@ Random.seed!(1234)
             for n in 0:L-1
                 @test_throws ErrorException transcode(ZstdDecompressor, compressed_data[1:n])
             end
+            @test CodecZstd.find_decompressed_size(compressed_data) == length(uncompressed_data)
             @test transcode(ZstdDecompressor, compressed_data) == uncompressed_data
         end
     end