[io] Compress buffers larger than 16 MiB in blocks of equal size (rev 2)

core/zip/inc/RZip.h

@@ -10,6 +10,9 @@
  *************************************************************************/
 #include "Compression.h"
 
+#include <cassert>
+#include <cstdint>
+
 /**
  * These are definitions of various free functions for the C-style compression routines in ROOT.
  */
@@ -39,4 +42,26 @@ extern "C" int R__unzip_header(int *srcsize, const unsigned char *src, int *tgts
 
 enum { kMAXZIPBUF = 0xffffff }; // 16 MB
 
+namespace ROOT::Internal {
+
+/// Number and size of the blocks in which a buffer is compressed.
+struct RZipBlocks {
+   int fNBlocks = 0;
+   int fBlockSize = 0;
+};
+
+/// Number of blocks and the size of the blocks in which a buffer of `nbytes` > 0 bytes is compressed.
+/// The buffer is cut into the fewest blocks of at most kMAXZIPBUF bytes, `1 + (nbytes - 1) / kMAXZIPBUF`, all
+/// of the same size but the last, which is smaller by less than the number of blocks. Full blocks followed by
+/// the remainder would end a buffer just above a multiple of kMAXZIPBUF in a block of a few bytes, which cannot
+/// be compressed, so that the whole buffer would be stored uncompressed.
+constexpr RZipBlocks GetZipBlocks(int nbytes)
+{
+   assert(nbytes > 0);
+   const std::uint64_t nblocks = 1 + (nbytes - 1) / kMAXZIPBUF;
+   return RZipBlocks{static_cast<int>(nblocks), static_cast<int>((nbytes + nblocks - 1) / nblocks)};
+}
+
+} // namespace ROOT::Internal
+
 #endif

io/io/src/TKey.cxx

@@ -262,16 +262,20 @@ TKey::TKey(const TObject *obj, const char *name, Int_t bufsize, TDirectory* moth
    Int_t cxlevel = GetFile() ? GetFile()->GetCompressionLevel() : 0;
    ROOT::RCompressionSetting::EAlgorithm::EValues cxAlgorithm = static_cast<ROOT::RCompressionSetting::EAlgorithm::EValues>(GetFile() ? GetFile()->GetCompressionAlgorithm() : 0);
    if (cxlevel > 0 && fObjlen > 256) {
-      Int_t nbuffers = 1 + (fObjlen - 1)/kMAXZIPBUF;
+      const auto blocks = ROOT::Internal::GetZipBlocks(fObjlen);
+      const Int_t nbuffers = blocks.fNBlocks;
+      const Int_t blockSize = blocks.fBlockSize;
       Int_t buflen = std::max(512,fKeylen + fObjlen + 9*nbuffers + 28); //add 28 bytes in case object is placed in a deleted gap
       fBuffer = new char[buflen];
       char *objbuf = fBufferRef->Buffer() + fKeylen;
       char *bufcur = &fBuffer[fKeylen];
       noutot = 0;
       nzip   = 0;
       for (Int_t i = 0; i < nbuffers; ++i) {
-         if (i == nbuffers - 1) bufmax = fObjlen - nzip;
-         else               bufmax = kMAXZIPBUF;
+         if (i == nbuffers - 1)
+            bufmax = fObjlen - nzip;
+         else
+            bufmax = blockSize;
          R__zipMultipleAlgorithm(cxlevel, &bufmax, objbuf, &bufmax, bufcur, &nout, cxAlgorithm);
          if (nout == 0 || nout >= fObjlen) { //this happens when the buffer cannot be compressed
             delete[] fBuffer;
@@ -284,8 +288,8 @@ TKey::TKey(const TObject *obj, const char *name, Int_t bufsize, TDirectory* moth
          }
          bufcur += nout;
          noutot += nout;
-         objbuf += kMAXZIPBUF;
-         nzip   += kMAXZIPBUF;
+         objbuf += blockSize;
+         nzip += blockSize;
       }
       Create(noutot);
       fBufferRef->SetBufferOffset(0);
@@ -357,16 +361,20 @@ TKey::TKey(const void *obj, const TClass *cl, const char *name, Int_t bufsize, T
    Int_t cxlevel = GetFile() ? GetFile()->GetCompressionLevel() : 0;
    ROOT::RCompressionSetting::EAlgorithm::EValues cxAlgorithm = static_cast<ROOT::RCompressionSetting::EAlgorithm::EValues>(GetFile() ? GetFile()->GetCompressionAlgorithm() : 0);
    if (cxlevel > 0 && fObjlen > 256) {
-      Int_t nbuffers = 1 + (fObjlen - 1)/kMAXZIPBUF;
+      const auto blocks = ROOT::Internal::GetZipBlocks(fObjlen);
+      const Int_t nbuffers = blocks.fNBlocks;
+      const Int_t blockSize = blocks.fBlockSize;
       Int_t buflen = std::max(512,fKeylen + fObjlen + 9*nbuffers + 28); //add 28 bytes in case object is placed in a deleted gap
       fBuffer = new char[buflen];
       char *objbuf = fBufferRef->Buffer() + fKeylen;
       char *bufcur = &fBuffer[fKeylen];
       noutot = 0;
       nzip   = 0;
       for (Int_t i = 0; i < nbuffers; ++i) {
-         if (i == nbuffers - 1) bufmax = fObjlen - nzip;
-         else               bufmax = kMAXZIPBUF;
+         if (i == nbuffers - 1)
+            bufmax = fObjlen - nzip;
+         else
+            bufmax = blockSize;
          R__zipMultipleAlgorithm(cxlevel, &bufmax, objbuf, &bufmax, bufcur, &nout, cxAlgorithm);
          if (nout == 0 || nout >= fObjlen) { //this happens when the buffer cannot be compressed
             delete[] fBuffer;
@@ -379,8 +387,8 @@ TKey::TKey(const void *obj, const TClass *cl, const char *name, Int_t bufsize, T
          }
          bufcur += nout;
          noutot += nout;
-         objbuf += kMAXZIPBUF;
-         nzip   += kMAXZIPBUF;
+         objbuf += blockSize;
+         nzip += blockSize;
       }
       Create(noutot);
       fBufferRef->SetBufferOffset(0);

io/io/test/ZipHeader.cxx

@@ -280,3 +280,96 @@ TEST(RZip, CorruptHeaderMessage)
 
    EXPECT_EQ(0, msg.Uncompress());
 }
+
+// Buffers one byte larger than the largest compression block (kMAXZIPBUF) are compressed in two blocks of about the
+// same size. With a full block followed by the remainder, the 1-byte trailing block could not be compressed and the
+// entire buffer was stored uncompressed.
+TEST(RZip, JustAboveBlockSizeTKey)
+{
+   TMemFile writableFile("memfile.root", "RECREATE");
+   writableFile.SetCompressionSettings(505);
+
+   // The streamed std::string is its characters plus a fixed overhead: measure it, then size the string such that
+   // the key's object is kMAXZIPBUF + 1 bytes long
+   std::string probe(1000, 'x');
+   writableFile.WriteObject(&probe, "probe");
+   int overhead = -1;
+   for (const auto &ki : writableFile.WalkTKeys()) {
+      if (ki.fKeyName == "probe")
+         overhead = ki.fObjLen - 1000;
+   }
+   ASSERT_GT(overhead, 0);
+   std::string large(kMAXZIPBUF + 1 - overhead, 'x');
+   writableFile.WriteObject(&large, "large");
+
+   bool found = false;
+   for (const auto &ki : writableFile.WalkTKeys()) {
+      if (ki.fKeyName != "large")
+         continue;
+      found = true;
+      EXPECT_EQ(kMAXZIPBUF + 1, ki.fObjLen);
+      EXPECT_LT(ki.fLen, ki.fObjLen); // compressed
+   }
+   EXPECT_TRUE(found);
+
+   auto readBack = writableFile.Get<std::string>("large");
+   ASSERT_TRUE(readBack);
+   EXPECT_EQ(large, *readBack);
+   delete readBack;
+}
+
+TEST(RZip, JustAboveBlockSizeTTree)
+{
+   TMemFile writableFile("memfile.root", "RECREATE");
+   writableFile.SetCompressionSettings(505);
+   auto tree = new TTree("t", "");
+   tree->SetAutoFlush(0);
+   Char_t val = 'x';
+   tree->Branch("val", &val, "val/B", 64 * 1024 * 1024);
+   const Long64_t nEntries = Long64_t(kMAXZIPBUF) + 1; // one basket of kMAXZIPBUF + 1 bytes
+   for (Long64_t i = 0; i < nEntries; ++i)
+      tree->Fill();
+   tree->Write();
+
+   int nBaskets = 0;
+   for (const auto &ki : writableFile.WalkTKeys()) {
+      if (ki.fClassName != "TBasket")
+         continue;
+      ++nBaskets;
+      EXPECT_EQ(kMAXZIPBUF + 1, ki.fObjLen);
+      EXPECT_LT(ki.fLen, ki.fObjLen); // compressed
+   }
+   EXPECT_EQ(1, nBaskets);
+
+   tree->SetBranchAddress("val", &val);
+   val = 0;
+   EXPECT_EQ(1, tree->GetEntry(nEntries - 1));
+   EXPECT_EQ('x', val);
+}
+
+TEST(RZip, JustAboveBlockSizeMessage)
+{
+   // The message length is the name plus a fixed overhead: measure it, then size the name such that the compressed
+   // part of the message (all but its 8-byte header) is kMAXZIPBUF + 1 bytes long
+   int overhead = 0;
+   {
+      TNamed probe;
+      probe.SetName(std::string(1000, 'x').c_str());
+      TMessage msg(kMESS_OBJECT | kMESS_ZIP);
+      msg.WriteObject(&probe);
+      overhead = msg.Length() - 2 * sizeof(UInt_t) - 1000;
+   }
+   ASSERT_GT(overhead, 0);
+
+   TNamed named;
+   named.SetName(std::string(kMAXZIPBUF + 1 - overhead, 'x').c_str());
+   TMessage msg(kMESS_OBJECT | kMESS_ZIP);
+   msg.SetCompressionSettings(505);
+   msg.WriteObject(&named);
+   EXPECT_EQ(kMAXZIPBUF + 1, msg.Length() - 2 * static_cast<int>(sizeof(UInt_t)));
+
+   EXPECT_EQ(0, msg.Compress());
+   EXPECT_TRUE(msg.CompBuffer());
+   EXPECT_LT(msg.CompLength(), msg.Length()); // compressed
+   EXPECT_EQ(0, msg.Uncompress());
+}

net/net/src/TMessage.cxx

@@ -359,7 +359,9 @@ Int_t TMessage::Compress()
 
    Int_t hdrlen   = 2*sizeof(UInt_t);
    Int_t messlen  = Length() - hdrlen;
-   Int_t nbuffers = 1 + (messlen - 1) / kMAXZIPBUF;
+   const auto blocks = ROOT::Internal::GetZipBlocks(messlen);
+   const Int_t nbuffers = blocks.fNBlocks;
+   const Int_t blockSize = blocks.fBlockSize;
    Int_t chdrlen  = 3*sizeof(UInt_t);   // compressed buffer header length
    Int_t buflen   = std::max(512, chdrlen + messlen + 9*nbuffers);
    fBufComp       = new char[buflen];
@@ -372,7 +374,7 @@ Int_t TMessage::Compress()
       if (i == nbuffers - 1)
          bufmax = messlen - nzip;
       else
-         bufmax = kMAXZIPBUF;
+         bufmax = blockSize;
       R__zipMultipleAlgorithm(compressionLevel, &bufmax, messbuf, &bufmax, bufcur, &nout,
                               static_cast<ROOT::RCompressionSetting::EAlgorithm::EValues>(compressionAlgorithm));
       if (nout == 0 || nout >= messlen) {
@@ -385,8 +387,8 @@ Int_t TMessage::Compress()
          return -1;
       }
       bufcur  += nout;
-      messbuf += kMAXZIPBUF;
-      nzip    += kMAXZIPBUF;
+      messbuf += blockSize;
+      nzip += blockSize;
    }
    fBufCompCur = bufcur;
    fCompPos    = fBufCur;

tree/ntuple/inc/ROOT/RNTupleZip.hxx

@@ -49,21 +49,23 @@ public:
       R__ASSERT(from != nullptr);
       R__ASSERT(to != nullptr);
       auto cxLevel = compression % 100;
-      if (cxLevel == 0) {
+      if (cxLevel == 0 || nbytes == 0) {
          memcpy(to, from, nbytes);
          return nbytes;
       }
 
       auto cxAlgorithm = static_cast<ROOT::RCompressionSetting::EAlgorithm::EValues>(compression / 100);
-      unsigned int nZipBlocks = 1 + (nbytes - 1) / kMAXZIPBUF;
+      const auto blocks = ROOT::Internal::GetZipBlocks(nbytes);
+      const unsigned int nZipBlocks = blocks.fNBlocks;
+      const int szBlock = blocks.fBlockSize;
       const char *source = static_cast<const char *>(from);
       int szTarget = nbytes;
       char *target = reinterpret_cast<char *>(to);
       int szOutBlock = 0;
       int szRemaining = nbytes;
       size_t szZipData = 0;
       for (unsigned int i = 0; i < nZipBlocks; ++i) {
-         int szSource = std::min(static_cast<int>(kMAXZIPBUF), szRemaining);
+         int szSource = std::min(szBlock, szRemaining);
          R__zipMultipleAlgorithm(cxLevel, &szSource, source, &szTarget, target, &szOutBlock, cxAlgorithm);
          R__ASSERT(szOutBlock >= 0);
          if ((szOutBlock == 0) || (szOutBlock >= szSource)) {

tree/ntuple/test/ntuple_zip.cxx

@@ -47,16 +47,35 @@ TEST(RNTupleZip, LargeWithOutputBuffer)
    auto unzipBuffer = MakeUninitArray<char>(N);
    std::string data(N, 'x');
 
-   /// Trailing byte cannot be compressed, entire buffer returns uncompressed
+   /// One byte more than a block: two blocks of about the same size, both compressed (no 1-byte trailing block,
+   /// which could not be compressed and would leave the entire buffer uncompressed)
    auto szZip = RNTupleCompressor::Zip(data.data(), kMAXZIPBUF + 1, 101, zipBuffer.get());
-   EXPECT_EQ(static_cast<unsigned int>(kMAXZIPBUF) + 1, szZip);
+   EXPECT_LT(szZip, static_cast<unsigned int>(kMAXZIPBUF) + 1);
+   RNTupleDecompressor::Unzip(zipBuffer.get(), szZip, kMAXZIPBUF + 1, unzipBuffer.get());
+   EXPECT_EQ(std::string_view(data.data(), kMAXZIPBUF + 1), std::string_view(unzipBuffer.get(), kMAXZIPBUF + 1));
 
    szZip = RNTupleCompressor::Zip(data.data(), data.length(), 101, zipBuffer.get());
    EXPECT_LT(szZip, N);
    RNTupleDecompressor::Unzip(zipBuffer.get(), szZip, N, unzipBuffer.get());
    EXPECT_EQ(data, std::string_view(unzipBuffer.get(), N));
 }
 
+TEST(RNTupleZip, LargeJustAboveBlocks)
+{
+   // Buffers just above one or two blocks: a page of 2^24 or 2^25 bytes must not be stored uncompressed
+   for (std::size_t n : {std::size_t(kMAXZIPBUF) + 1, std::size_t(kMAXZIPBUF) + 2, std::size_t(1) << 25}) {
+      auto zipBuffer = MakeUninitArray<unsigned char>(n);
+      auto unzipBuffer = MakeUninitArray<char>(n);
+      std::string data(n, 'x');
+      for (int compression : {101, 207, 404, 505}) {
+         auto szZip = RNTupleCompressor::Zip(data.data(), n, compression, zipBuffer.get());
+         EXPECT_LT(szZip, n) << "size " << n << ", compression " << compression;
+         RNTupleDecompressor::Unzip(zipBuffer.get(), szZip, n, unzipBuffer.get());
+         EXPECT_EQ(data, std::string_view(unzipBuffer.get(), n));
+      }
+   }
+}
+
 TEST(RNTupleZip, CorruptedInput)
 {
    std::string data = "xxxxxxxxxxxxxxxxxxxxxxxx";

tree/tree/src/TBasket.cxx

@@ -1297,8 +1297,10 @@ Int_t TBasket::WriteBuffer()
    ROOT::RCompressionSetting::EAlgorithm::EValues cxAlgorithm = static_cast<ROOT::RCompressionSetting::EAlgorithm::EValues>(fBranch->GetCompressionAlgorithm());
    if (cxAlgorithm == ROOT::RCompressionSetting::EAlgorithm::kInherit)
       cxAlgorithm = static_cast<ROOT::RCompressionSetting::EAlgorithm::EValues>(file->GetCompressionAlgorithm());
-   if (cxlevel > 0) {
-      Int_t nbuffers = 1 + (fObjlen - 1) / kMAXZIPBUF;
+   if (cxlevel > 0 && fObjlen > 0) {
+      const auto blocks = ROOT::Internal::GetZipBlocks(fObjlen);
+      const Int_t nbuffers = blocks.fNBlocks;
+      const Int_t blockSize = blocks.fBlockSize;
       Int_t buflen = fKeylen + fObjlen + 9 * nbuffers + 28; //add 28 bytes in case object is placed in a deleted gap
       InitializeCompressedBuffer(buflen, file);
       if (!fCompressedBufferRef) {
@@ -1312,8 +1314,10 @@ Int_t TBasket::WriteBuffer()
       noutot = 0;
       nzip   = 0;
       for (Int_t i = 0; i < nbuffers; ++i) {
-         if (i == nbuffers - 1) bufmax = fObjlen - nzip;
-         else bufmax = kMAXZIPBUF;
+         if (i == nbuffers - 1)
+            bufmax = fObjlen - nzip;
+         else
+            bufmax = blockSize;
          // Compress the buffer.  Note that we allow multiple TBasket compressions to occur at once
          // for a given TFile: that's because the compression buffer when we use IMT is no longer
          // shared amongst several threads.
@@ -1348,8 +1352,8 @@ Int_t TBasket::WriteBuffer()
          }
          bufcur += nout;
          noutot += nout;
-         objbuf += kMAXZIPBUF;
-         nzip   += kMAXZIPBUF;
+         objbuf += blockSize;
+         nzip += blockSize;
       }
       nout = noutot;
       Create(noutot,file);