RecordFileWriter.h

/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
 
#pragma once
 
#include <atomic>
#include <functional>
#include <memory>
#include <set>
#include <thread>
 
#include <vrs/Compressor.h>
#include <vrs/DiskFile.h>
#include <vrs/FileFormat.h>
#include <vrs/IndexRecord.h>
#include <vrs/NewChunkHandler.h>
#include <vrs/Record.h>
#include <vrs/Recordable.h>
#include <vrs/WriteFileHandler.h>
 
namespace vrs {
 
namespace test {
struct RecordFileWriterTester;
} // namespace test
 
using std::function;
using std::map;
using std::pair;
using std::set;
using std::string;
 
class Record;
class Recordable;
 
namespace RecordFileWriter_ {
struct WriterThreadData;
struct PurgeThreadData;
struct CompressionThreadsData;
struct RecordWriterData;
} // namespace RecordFileWriter_
 
enum class ThreadRole { Writer, Purge, Compression };
 
using InitCreatedThreadCallback = std::function<void(std::thread&, ThreadRole, int)>;
 
class RecordFileWriter {
 public:
  RecordFileWriter();
  RecordFileWriter(const RecordFileWriter&) = delete;
  RecordFileWriter(RecordFileWriter&&) = delete;
  virtual ~RecordFileWriter() noexcept;
 
  RecordFileWriter& operator=(const RecordFileWriter&) = delete;
  RecordFileWriter& operator=(RecordFileWriter&&) = delete;
 
  void addRecordable(Recordable* recordable);
 
  vector<Recordable*> getRecordables() const;
 
  void setCompressionThreadPoolSize(size_t size = kMaxThreadPoolSizeForHW);
  static constexpr size_t kMaxThreadPoolSizeForHW = UINT32_MAX; // one-per-CPU core
 
  void setInitCreatedThreadCallback(const InitCreatedThreadCallback& initCreatedThreadCallback);
 
  int writeToFile(const string& filePath);
 
  void purgeOldRecords(double maxTimestamp, bool recycleBuffers = true);
 
  virtual int createFileAsync(const string& filePath);
 
  int createChunkedFile(
      const string& filePath,
      size_t maxChunkSizeMB = 0,
      unique_ptr<NewChunkHandler>&& chunkHandler = nullptr);
 
  void setMaxChunkSizeMB(size_t maxChunkSizeMB);
 
  int preallocateIndex(unique_ptr<deque<IndexRecord::DiskRecordInfo>> preliminaryIndex);
 
  int writeRecordsAsync(double maxTimestamp);
 
  int autoWriteRecordsAsync(const function<double()>& maxTimestampProvider, double delay);
 
  int autoPurgeRecords(const function<double()>& maxTimestampProvider, double delay);
 
  bool isWriting() const {
    return file_->isOpened();
  }
 
  void trackBackgroundThreadQueueByteSize();
  uint64_t getBackgroundThreadQueueByteSize();
 
  int closeFileAsync();
 
  virtual int waitForFileClosed();
 
  void setTag(const string& tagName, const string& tagValue);
  void addTags(const map<string, string>& newTags);
  const map<string, string>& getTags() const {
    return fileTags_;
  }
 
  int setWriteFileHandler(unique_ptr<WriteFileHandler> writeFileHandler);
 
  void backgroundWriterThreadActivity();
  void backgroundPurgeThreadActivity();
 
  struct SortRecord {
    SortRecord(Record* recordIn, StreamId streamIdIn) : record(recordIn), streamId(streamIdIn) {}
 
    bool operator<(const SortRecord& rhs) const {
      return this->record->getTimestamp() < rhs.record->getTimestamp() ||
          (this->record->getTimestamp() <= rhs.record->getTimestamp() &&
           (this->streamId < rhs.streamId ||
            // Records have a unique creation order within a particular device
            (this->streamId == rhs.streamId &&
             this->record->getCreationOrder() < rhs.record->getCreationOrder())));
    }
 
    Record* record;
    StreamId streamId;
  };
  using RecordBatch = vector<pair<StreamId, list<Record*>>>;
  using RecordBatches = vector<unique_ptr<RecordBatch>>;
  using SortedRecords = deque<SortRecord>;
 
 private:
  uint64_t collectOldRecords(RecordBatch& batch, double maxTimestamp); 
  bool autoCollectRecords(bool checkTime); 
  static uint64_t addRecordBatchesToSortedRecords(
      const RecordBatches& batch,
      SortedRecords& inOutSortedRecords); 
  int createFileAsync(const string& filePath, bool splitHead); 
  int createFile(const string& filePath, bool splitHead); 
  int writeRecords(SortedRecords& records, int lastError); 
  int writeRecordsSingleThread(SortedRecords& records, int lastError); 
  int writeRecordsMultiThread(
      RecordFileWriter_::CompressionThreadsData& compressionThreadsData,
      SortedRecords& records,
      int lastError); 
  void writeOneRecord(
      RecordFileWriter_::RecordWriterData& rwd,
      Record* record,
      StreamId id,
      Compressor& compressor,
      uint32_t compressedSize);
  int completeAndCloseFile(); 
 
  void backgroundWriteCollectedRecord(); 
  bool addRecordsReadyToWrite(SortedRecords& inOutRecordsToWrite); 
 
  set<Recordable*> recordables_;
  mutable std::mutex recordablesMutex_;
 
 protected:
  unique_ptr<WriteFileHandler> file_;
  uint64_t maxChunkSize_{};
  unique_ptr<NewChunkHandler> newChunkHandler_;
  FileFormat::FileHeader fileHeader_;
  uint32_t lastRecordSize_{};
  bool sortRecords_{true}; // should records be sorted? Don't only for best IO throughput.
  bool skipFinalizeIndexRecords_ = false; // for unit testing only!
  unique_ptr<deque<IndexRecord::DiskRecordInfo>> preliminaryIndex_;
  IndexRecord::Writer indexRecordWriter_;
  map<string, string> fileTags_;
  size_t compressionThreadPoolSize_{};
 
  RecordFileWriter_::WriterThreadData* writerThreadData_;
  std::unique_ptr<std::atomic<uint64_t>> queueByteSize_; // background thread's queue byte size
 
  RecordFileWriter_::PurgeThreadData* purgeThreadData_;
 
  InitCreatedThreadCallback initCreatedThreadCallback_;
 
  friend struct ::vrs::test::RecordFileWriterTester; // for tests ONLY
};
 
} // namespace vrs