writer.h

#ifndef SINGA_IO_WRITER_H_
#define SINGA_IO_WRITER_H_

#include <cstring>
#include <fstream>
#include <string>
#include "singa/singa_config.h"

#ifdef USE_LMDB
#include <lmdb.h>
#include <sys/stat.h>
#include <vector>
#endif  // USE_LMDB

namespace singa {
namespace io {

using std::string;
enum Mode { kCreate, kAppend };

class Writer {
 public:
  virtual ~Writer() {}

  virtual bool Open(const std::string &path, Mode mode) = 0;

  virtual void Close() = 0;

  virtual bool Write(const std::string &key, const std::string &value) = 0;

  virtual void Flush() = 0;
};

class BinFileWriter : public Writer {
 public:
  ~BinFileWriter() { Close(); }
  bool Open(const std::string &path, Mode mode) override;
  bool Open(const std::string &path, Mode mode, int capacity);
  void Close() override;
  bool Write(const std::string &key, const std::string &value) override;
  void Flush() override;
  inline std::string path() { return path_; }

 protected:
  bool OpenFile();

 private:
  std::string path_ = "";
  Mode mode_;
  std::ofstream fdat_;
  char *buf_ = nullptr;
  int capacity_ = 10485760;
  int bufsize_ = 0;
  const char kMagicWord[2] = {'s', 'g'};
};

class TextFileWriter : public Writer {
 public:
  ~TextFileWriter() { Close(); }
  bool Open(const std::string &path, Mode mode) override;
  void Close() override;
  bool Write(const std::string &key, const std::string &value) override;
  void Flush() override;
  inline std::string path() { return path_; }

 private:
  std::string path_ = "";
  Mode mode_;
  std::ofstream fdat_;
};

#ifdef USE_LMDB
class LMDBWriter : public Writer {
 public:
  ~LMDBWriter() { Close(); }
  bool Open(const std::string &path, Mode mode) override;
  void Close() override;
  bool Write(const std::string &key, const std::string &value) override;
  void Flush() override;
  inline std::string path() { return path_; }

 protected:
  void DoubleMapSize();
  inline void MDB_CHECK(int mdb_status);

 private:
  std::string path_ = "";
  Mode mode_;
  MDB_env *mdb_env_ = nullptr;
  std::vector<string> keys, values;
};
#endif  // USE_LMDB

}  // namespace io
}  // namespace singa

#endif  // SINGA_IO_WRITER_H_