Unbuffered variant of file qstorage (#4194)

Author: vitstn

ydb/library/yql/core/qplayer/storage/file/ut/yql_qstorage_file_ut.cpp

@@ -6,6 +6,22 @@
 
 using namespace NYql;
 
-Y_UNIT_TEST_SUITE(TQStorageFileTests) {
-    GENERATE_TESTS(MakeFileQStorage)
+IQStoragePtr MakeBufferedFileQStorage() {
+    TFileQStorageSettings settings;
+    settings.BufferUntilCommit = true;
+    return MakeFileQStorage({}, settings);
+}
+
+IQStoragePtr MakeUnbufferedFileQStorage() {
+    TFileQStorageSettings settings;
+    settings.BufferUntilCommit = false;
+    return MakeFileQStorage({}, settings);
+}
+
+Y_UNIT_TEST_SUITE(TQStorageBufferedFileTests) {
+    GENERATE_TESTS(MakeBufferedFileQStorage)
+}
+
+Y_UNIT_TEST_SUITE(TQStorageUnbufferedFileTests) {
+    GENERATE_TESTS(MakeUnbufferedFileQStorage)
 }

ydb/library/yql/core/qplayer/storage/file/yql_qstorage_file.cpp

@@ -14,10 +14,44 @@ namespace NYql {
 
 namespace {
 
-class TWriter : public IQWriter {
-public:
-    TWriter(TFsPath& path)
+class TWriterBase : public IQWriter {
+protected:
+    TWriterBase(TFsPath& path, TInstant writtenAt)
         : Path_(path)
+        , WrittenAt_(writtenAt)
+    {}
+
+protected:
+    void WriteIndex(ui64 totalItems, ui64 totalBytes, ui64 checksum) const {
+        TFileOutput indexFile(Path_.GetPath() + ".idx.tmp");
+        indexFile.Write(&WrittenAt_, sizeof(WrittenAt_));
+        indexFile.Write(&totalItems, sizeof(totalItems));
+        indexFile.Write(&totalBytes, sizeof(totalBytes));
+        indexFile.Write(&checksum, sizeof(checksum));
+        indexFile.Finish();
+        if (!NFs::Rename(Path_.GetPath() + ".idx.tmp", Path_.GetPath() + ".idx")) {
+            throw yexception() << "can not rename: " << LastSystemErrorText();
+        }
+    }
+
+    static void SaveString(TFileOutput& file, const TString& str, ui64& totalBytes, ui64& checksum) {
+        ui32 length = str.Size();
+        checksum = crc64(&length, sizeof(length), checksum);
+        file.Write(&length, sizeof(length));
+        checksum = crc64(str.Data(), length, checksum);
+        file.Write(str.Data(), length);
+        totalBytes += length;
+    }
+
+protected:
+    const TFsPath Path_;
+    const TInstant WrittenAt_;
+};
+
+class TBufferedWriter : public TWriterBase {
+public:
+    TBufferedWriter(TFsPath& path, TInstant writtenAt)
+        : TWriterBase(path, writtenAt)
         , Storage_(MakeMemoryQStorage())
         , Writer_(Storage_->MakeWriter("", {}))
     {
@@ -36,6 +70,7 @@ class TWriter : public IQWriter {
 private:
     void SaveFile(const IQIteratorPtr& iterator) {
         TFileOutput dataFile(Path_.GetPath() + ".dat");
+        dataFile.Write(&WrittenAt_, sizeof(WrittenAt_));
         ui64 totalItems = 0;
         ui64 totalBytes = 0;
         ui64 checksum = 0;
@@ -52,45 +87,77 @@ class TWriter : public IQWriter {
         }
 
         dataFile.Finish();
-        TFileOutput indexFile(Path_.GetPath() + ".idx.tmp");
-        indexFile.Write(&totalItems, sizeof(totalItems));
-        indexFile.Write(&totalBytes, sizeof(totalBytes));
-        indexFile.Write(&checksum, sizeof(checksum));
-        if (!NFs::Rename(Path_.GetPath() + ".idx.tmp", Path_.GetPath() + ".idx")) {
-            throw yexception() << "can not rename: " << LastSystemErrorText();
+        WriteIndex(totalItems, totalBytes, checksum);
+    }
+
+private:
+    const IQStoragePtr Storage_;
+    const IQWriterPtr Writer_;
+};
+
+class TUnbufferedWriter : public TWriterBase {
+public:
+    TUnbufferedWriter(TFsPath& path, TInstant writtenAt)
+        : TWriterBase(path, writtenAt)
+        , DataFile_(Path_.GetPath() + ".dat")
+    {
+        DataFile_.Write(&WrittenAt_, sizeof(WrittenAt_));
+    }
+
+    NThreading::TFuture<void> Put(const TQItemKey& key, const TString& value) final {
+        with_lock(Mutex_) {
+            Y_ENSURE(!Committed_);
+            if (Keys_.emplace(key).second) {
+                SaveString(DataFile_, key.Component, TotalBytes_, Checksum_);
+                SaveString(DataFile_, key.Label, TotalBytes_, Checksum_);
+                SaveString(DataFile_, value, TotalBytes_, Checksum_);
+                ++TotalItems_;
+            }
+
+            return NThreading::MakeFuture();
         }
     }
 
-    void SaveString(TFileOutput& file, const TString& str, ui64& totalBytes, ui64& checksum) {
-        ui32 length = str.Size();
-        checksum = crc64(&length, sizeof(length), checksum);
-        file.Write(&length, sizeof(length));
-        checksum = crc64(str.Data(), length, checksum);
-        file.Write(str.Data(), length);
-        totalBytes += length;
+    NThreading::TFuture<void> Commit() final {
+        with_lock(Mutex_) {
+            Y_ENSURE(!Committed_);
+            Committed_ = true;
+            DataFile_.Finish();
+            WriteIndex(TotalItems_, TotalBytes_, Checksum_);
+            return NThreading::MakeFuture();
+        }
     }
 
 private:
-    const TFsPath Path_;
-    const IQStoragePtr Storage_;
-    const IQWriterPtr Writer_;
+    TMutex Mutex_;
+    TFileOutput DataFile_;
+    ui64 TotalItems_ = 0;
+    ui64 TotalBytes_ = 0;
+    ui64 Checksum_ = 0;
+    THashSet<TQItemKey> Keys_;
+    bool Committed_ = false;
 };
 
 class TStorage : public IQStorage {
 public:
-    TStorage(const TString& folder)
+    TStorage(const TString& folder, const TFileQStorageSettings& settings)
         : Folder_(folder)
+        , Settings_(settings)
     {
         if (!Folder_.IsDefined()) {
             TmpDir_.ConstructInPlace();
             Folder_ = TmpDir_->Path();
         }
     }
 
-    IQWriterPtr MakeWriter(const TString& operationId, const TQWriterSettings& settings) const final {
-        Y_UNUSED(settings);
+    IQWriterPtr MakeWriter(const TString& operationId, const TQWriterSettings& writerSettings) const final {
         auto opPath = Folder_ / operationId;
-        return std::make_shared<TWriter>(opPath);
+        auto writtenAt = writerSettings.WrittenAt.GetOrElse(Now());
+        if (Settings_.BufferUntilCommit) {
+            return std::make_shared<TBufferedWriter>(opPath, writtenAt);
+        } else {
+            return std::make_shared<TUnbufferedWriter>(opPath, writtenAt);
+        }
     }
 
     IQReaderPtr MakeReader(const TString& operationId, const TQReaderSettings& settings) const final {
@@ -115,14 +182,19 @@ class TStorage : public IQStorage {
 
         auto writer = memory->MakeWriter("", {});
         TFileInput indexFile(indexPath.GetPath());
+        TInstant indexWrittenAt;
         ui64 totalItems, loadedTotalBytes, loadedChecksum;
+        indexFile.LoadOrFail(&indexWrittenAt, sizeof(indexWrittenAt));
         indexFile.LoadOrFail(&totalItems, sizeof(totalItems));
         indexFile.LoadOrFail(&loadedTotalBytes, sizeof(loadedTotalBytes));
         indexFile.LoadOrFail(&loadedChecksum, sizeof(loadedChecksum));
         char dummy;
         Y_ENSURE(!indexFile.ReadChar(dummy));
         const TFsPath& dataPath = Folder_ / (operationId + ".dat");
         TFileInput dataFile(dataPath.GetPath());
+        TInstant dataWrittenAt;
+        dataFile.LoadOrFail(&dataWrittenAt, sizeof(dataWrittenAt));
+        Y_ENSURE(indexWrittenAt == dataWrittenAt);
         ui64 totalBytes = 0, checksum = 0;
         for (ui64 i = 0; i < totalItems; ++i) {
             TQItemKey key;
@@ -160,12 +232,13 @@ class TStorage : public IQStorage {
 private:
     TMaybe<TTempDir> TmpDir_;
     TFsPath Folder_;
+    const TFileQStorageSettings Settings_;
 };
 
 }
 
-IQStoragePtr MakeFileQStorage(const TString& folder) {
-    return std::make_shared<TStorage>(folder);
+IQStoragePtr MakeFileQStorage(const TString& folder, const TFileQStorageSettings& settings) {
+    return std::make_shared<TStorage>(folder, settings);
 }
 
 };

ydb/library/yql/core/qplayer/storage/file/yql_qstorage_file.h

@@ -3,6 +3,10 @@
 
 namespace NYql {
 
-IQStoragePtr MakeFileQStorage(const TString& folder = {});
+struct TFileQStorageSettings {
+    bool BufferUntilCommit = true;
+};
+
+IQStoragePtr MakeFileQStorage(const TString& folder = {}, const TFileQStorageSettings& settings = {});
 
 };

ydb/library/yql/core/qplayer/storage/memory/yql_qstorage_memory.cpp

@@ -53,7 +53,7 @@ class TWriter : public IQWriter {
     NThreading::TFuture<void> Put(const TQItemKey& key, const TString& value) final {
         with_lock(Operation_->Mutex) {
             Y_ENSURE(!Operation_->Committed);
-            (*Operation_->WriteMap)[key] = value;
+            Operation_->WriteMap->emplace(key, value);
             return NThreading::MakeFuture();
         }
     }