Cache context in thread_local rather than class member#48187

Author: Andreas Ringlstetter

cpp/src/arrow/util/compression_zstd.cc

@@ -202,11 +202,12 @@ class ZSTDCodec : public Codec {
       DCHECK_EQ(output_buffer_len, 0);
       output_buffer = &empty_buffer;
     }
-    if (!decompression_context_) {
-      decompression_context_ =
-          decltype(decompression_context_)(ZSTD_createDCtx(), ZSTDContextDeleter{});
-    }
-    size_t ret = ZSTD_decompressDCtx(decompression_context_.get(), output_buffer,
+    // Decompression context for ZSTD contains several large heap allocations.
+    // This method is additionally used in a free-threaded context so caching in a class
+    // member is not possible.
+    thread_local std::unique_ptr<ZSTD_DCtx, ZSTDContextDeleter> decompression_context{
+        ZSTD_createDCtx(), ZSTDContextDeleter{}};
+    size_t ret = ZSTD_decompressDCtx(decompression_context.get(), output_buffer,
                                      static_cast<size_t>(output_buffer_len), input,
                                      static_cast<size_t>(input_len));
     if (ZSTD_isError(ret)) {
@@ -226,18 +227,14 @@ class ZSTDCodec : public Codec {
 
   Result<int64_t> Compress(int64_t input_len, const uint8_t* input,
                            int64_t output_buffer_len, uint8_t* output_buffer) override {
-    if (!compression_context_) {
-      compression_context_ =
-          decltype(compression_context_)(ZSTD_createCCtx(), ZSTDContextDeleter{});
-      size_t ret = ZSTD_CCtx_setParameter(compression_context_.get(),
-                                          ZSTD_c_compressionLevel, compression_level_);
-      if (ZSTD_isError(ret)) {
-        return ZSTDError(ret, "Setting ZSTD compression level failed: ");
-      }
-    }
-    size_t ret = ZSTD_compress2(compression_context_.get(), output_buffer,
-                                static_cast<size_t>(output_buffer_len), input,
-                                static_cast<size_t>(input_len));
+    // Compression context for ZSTD contains several large heap allocations.
+    // This method is additionally used in a free-threaded context so caching in a class
+    // member is not possible.
+    thread_local std::unique_ptr<ZSTD_CCtx, ZSTDContextDeleter> compression_context{
+        ZSTD_createCCtx(), ZSTDContextDeleter{}};
+    size_t ret = ZSTD_compressCCtx(compression_context.get(), output_buffer,
+                                   static_cast<size_t>(output_buffer_len), input,
+                                   static_cast<size_t>(input_len), compression_level_);
     if (ZSTD_isError(ret)) {
       return ZSTDError(ret, "ZSTD compression failed: ");
     }
@@ -265,8 +262,6 @@ class ZSTDCodec : public Codec {
 
  private:
   const int compression_level_;
-  std::unique_ptr<ZSTD_DCtx, ZSTDContextDeleter> decompression_context_;
-  std::unique_ptr<ZSTD_CCtx, ZSTDContextDeleter> compression_context_;
 };
 
 }  // namespace