Add QWriter::Close()

commit_hash:43d2f04e77089ac19624b5ba9db9192446c23d0e

Author: MrLolthe1st

yql/essentials/core/qplayer/storage/file/yql_qstorage_file.cpp

@@ -7,8 +7,8 @@
 #include <util/folder/tempdir.h>
 #include <util/generic/hash_set.h>
 #include <util/system/fs.h>
-#include <util/system/mutex.h>
 #include <util/stream/file.h>
+#include <util/system/mutex.h>
 
 namespace NYql {
 
@@ -166,6 +166,12 @@ class TUnbufferedWriter : public TWriterBase {
         }
     }
 
+    void Close() override final {
+        with_lock(Mutex_) {
+            DataFile_.Clear();
+        }
+    }
+
 private:
     const TQWriterSettings Settings_;
     const bool AlwaysFlushIndex_;
@@ -195,9 +201,9 @@ class TStorage : public IQStorage {
         auto opPath = Folder_ / operationId;
         auto writtenAt = writerSettings.WrittenAt.GetOrElse(Now());
         if (Settings_.BufferUntilCommit) {
-            return std::make_shared<TBufferedWriter>(opPath, writtenAt, writerSettings);
+            return MakeCloseAwareWriterDecorator(std::make_shared<TBufferedWriter>(opPath, writtenAt, writerSettings));
         } else {
-            return std::make_shared<TUnbufferedWriter>(opPath, writtenAt, writerSettings, Settings_.AlwaysFlushIndex);
+            return MakeCloseAwareWriterDecorator(std::make_shared<TUnbufferedWriter>(opPath, writtenAt, writerSettings, Settings_.AlwaysFlushIndex));
         }
     }
 

yql/essentials/core/qplayer/storage/interface/yql_qstorage.cpp

@@ -1 +1,37 @@
 #include "yql_qstorage.h"
+
+namespace NYql {
+class TQWriterDecorator : public IQWriter {
+    public:
+    TQWriterDecorator(IQWriterPtr&& underlying) : Underlying_(std::move(underlying)) {}
+    NThreading::TFuture<void> Put(const TQItemKey& key, const TString& value) override final {
+        if (Closed_) {
+            return NThreading::MakeFuture();
+        }
+        return Underlying_->Put(key, value);
+    }
+
+    NThreading::TFuture<void> Commit() override final {
+        if (Closed_) {
+            throw yexception() << "QWriter closed";
+        }
+        return Underlying_->Commit();
+    }
+
+    // Close all used files, doesn't commit anything
+    void Close() override final {
+        bool expected = false;
+        if (Closed_.compare_exchange_strong(expected, true)) {
+            Underlying_ = {};
+        }
+    }
+private:
+    IQWriterPtr Underlying_;
+    std::atomic<bool> Closed_ = false;
+};
+
+IQWriterPtr MakeCloseAwareWriterDecorator(IQWriterPtr&& rhs) {
+    return std::make_shared<TQWriterDecorator>(std::move(rhs));
+}
+
+}

yql/essentials/core/qplayer/storage/interface/yql_qstorage.h

@@ -48,10 +48,10 @@ using IQReaderPtr = std::shared_ptr<IQReader>;
 class IQWriter {
 public:
     virtual ~IQWriter() = default;
-
     virtual NThreading::TFuture<void> Put(const TQItemKey& key, const TString& value) = 0;
     // Commmit should be called at most once, no more Put are allowed after it
     virtual NThreading::TFuture<void> Commit() = 0;
+    virtual void Close() {};
 };
 
 using IQWriterPtr = std::shared_ptr<IQWriter>;
@@ -134,6 +134,7 @@ class TQContext {
     IQWriterPtr Writer_;
 };
 
+IQWriterPtr MakeCloseAwareWriterDecorator(IQWriterPtr&& rhs);
 }
 
 template <>

yql/essentials/core/qplayer/storage/memory/yql_qstorage_memory.cpp

@@ -129,7 +129,7 @@ class TStorage : public IQStorage {
     }
 
     IQWriterPtr MakeWriter(const TString& operationId, const TQWriterSettings& writerSettings) const final {
-        return std::make_shared<TWriter>(GetOperation(operationId, true), writerSettings);
+        return MakeCloseAwareWriterDecorator(std::make_shared<TWriter>(GetOperation(operationId, true), writerSettings));
     }
 
     IQIteratorPtr MakeIterator(const TString& operationId, const TQIteratorSettings& iteratorSettings) const final {

yql/essentials/core/qplayer/storage/ydb/yql_qstorage_ydb.cpp

@@ -218,7 +218,7 @@ class TStorage : public IQStorage {
     }
 
     IQWriterPtr MakeWriter(const TString& operationId, const TQWriterSettings& settings) const final {
-        return std::make_shared<TWriter>(Settings_, operationId, settings);
+        return MakeCloseAwareWriterDecorator(std::make_shared<TWriter>(Settings_, operationId, settings));
     }
 
     IQReaderPtr MakeReader(const TString& operationId, const TQReaderSettings& settings) const final {
@@ -238,7 +238,7 @@ class TStorage : public IQStorage {
     void LoadTable(const TString& operationId, const IQStoragePtr& memory) const {
         auto driver = MakeDriver(Settings_);
         NYdb::NTable::TTableClient tableClient(driver);
-        
+
         auto operationsTable  = Settings_.TablesPrefix + "operations";
         auto fullOperationId = Settings_.OperationIdPrefix + operationId;
 
@@ -287,7 +287,7 @@ class TStorage : public IQStorage {
         TString blobTable = Settings_.Database + "/" + Settings_.TablesPrefix + "blobs";
 
         const auto maxBatchSize = Settings_.MaxBatchSize.GetOrElse(DefaultMaxBatchSize);
-        auto rtResult = tableClient.RetryOperationSync([&tableIter, maxBatchSize, blobTable, 
+        auto rtResult = tableClient.RetryOperationSync([&tableIter, maxBatchSize, blobTable,
             fullOperationId, writtenAt, loadedTotalItems](NYdb::NTable::TSession session) {
             auto key1 = NYdb::TValueBuilder()
                 .BeginTuple()
@@ -321,7 +321,7 @@ class TStorage : public IQStorage {
             if (res.IsSuccess()) {
                 tableIter = res;
             }
-            
+
             return res;
         }, readRetrySettings);
         ThrowOnError(rtResult);