FairMQResizableBuffer.h

Go to the documentation of this file.

// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
// All rights not expressly granted are reserved.
//
// This software is distributed under the terms of the GNU General Public
// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
//
// In applying this license CERN does not waive the privileges and immunities
// granted to it by virtue of its status as an Intergovernmental Organization
// or submit itself to any jurisdiction.
 
#ifndef O2_FRAMEWORK_FAIRMQRESIZABLEBUFFER_H_
#define O2_FRAMEWORK_FAIRMQRESIZABLEBUFFER_H_
 
#include <memory>
#include <functional>
#include <arrow/buffer.h>
#include "arrow/io/interfaces.h"
#include "arrow/status.h"
#include "arrow/util/future.h"
 
#include <fairmq/FwdDecls.h>
 
namespace o2::framework
{
 
class FairMQOutputStream : public arrow::io::OutputStream
{
 public:
  explicit FairMQOutputStream(const std::shared_ptr<arrow::ResizableBuffer>& buffer);
 
  static arrow::Result<std::shared_ptr<FairMQOutputStream>> Create(
    int64_t initial_capacity = 4096, arrow::MemoryPool* pool = arrow::default_memory_pool());
 
  // By the time we call the destructor, the contents
  // of the buffer are already moved to fairmq
  // for being sent.
  ~FairMQOutputStream() override = default;
 
  // Implement the OutputStream interface
 
  arrow::Status Close() override;
  [[nodiscard]] bool closed() const override;
  [[nodiscard]] arrow::Result<int64_t> Tell() const override;
  arrow::Status Write(const void* data, int64_t nbytes) override;
 
  using OutputStream::Write;
 
  arrow::Result<std::shared_ptr<arrow::Buffer>> Finish();
 
  arrow::Status Reset(int64_t initial_capacity = 1024, arrow::MemoryPool* pool = arrow::default_memory_pool());
 
  [[nodiscard]] int64_t capacity() const { return capacity_; }
 
 private:
  FairMQOutputStream();
 
  // Ensures there is sufficient space available to write nbytes
  arrow::Status Reserve(int64_t nbytes);
 
  std::shared_ptr<arrow::ResizableBuffer> buffer_;
  bool is_open_;
  int64_t capacity_;
  int64_t position_;
  uint8_t* mutable_data_;
};
 
class FairMQResizableBuffer : public ::arrow::ResizableBuffer
{
 public:
  using Creator = std::function<std::unique_ptr<fair::mq::Message>(size_t)>;
 
  FairMQResizableBuffer(Creator);
  ~FairMQResizableBuffer() override;
 
  arrow::Status Resize(const int64_t new_size, bool shrink_to_fit) override;
  arrow::Status Reserve(const int64_t capacity) override;
 
  std::unique_ptr<fair::mq::Message> Finalise();
 
 private:
  std::unique_ptr<fair::mq::Message> mMessage;
  int64_t mSize;
  Creator mCreator;
};
 
} // namespace o2::framework
 
#endif // O2_FRAMEWORK_FAIRMQRESIZABLEBUFFER_H_