// Copyright 2019 DeepMind Technologies Limited.

//

// 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.

#ifndef REVERB_CC_PLATFORM_TFRECORD_CHECKPOINTER_H_

#define REVERB_CC_PLATFORM_TFRECORD_CHECKPOINTER_H_

#include <memory>

#include <string>

#include <vector>

#include "absl/status/status.h"

#include "absl/strings/string_view.h"

#include "absl/types/optional.h"

#include "reverb/cc/checkpointing/interface.h"

#include "reverb/cc/chunk_store.h"

#include "reverb/cc/table.h"

namespace deepmind {

namespace reverb {

// Generates and stores proto checkpoints of PriorityTables and ChunkStore data

// to a directory inside the top level `root_dir`.

//

// A set of `Table` constitutes the bases for a checkpoint. When `Save` is

// called state of each `Table` is encoded into a PriorityTableCheckpoint. The

// proto contains the state and initialization options of the table itself and

// all its dependencies (RateLimiter, KeyDistribution etc) but does not include

// the PrioritizedItem or Chunk data. Instead two containers are created. The

// first is an ordered vector of PrioritizedItem copies (in ascending order of

// `inserted_at`). The second contains a shared_ptr to every referenced

// ChunkStore::Chunk is attached (which ensures that all data remains for the

// complete duration of the checkpointing operation).

//

// To avoid duplicating data, the union of the referenced chunks are

// deduplicated before being stored to disk. The stored checkpoint has the

// following format:

//

// <root_dir>/

// <timestamp of the checkpoint>/

// tables.tfrecord

// items.tfrecord

// chunks.tfrecord

// DONE

//

// DONE an empty file written once the checkpoint has been successfully written.

// If DONE does not exist then the checkpoint is in process of being written or

// the operation was unexpectedly interrupted and the data should be considered

// corrupt.

//

// The most recent checkpoint can therefore be inferred from the name of the

// directories within `root_dir`.

//

// If `group` is nonempty then the directory containing the checkpoint will be

// created with `group` as group.

//

// The optional class member `fallback_checkpoint_path` can be set to specify

// a checkpoint to be reloaded when no checkpoints can be found in `root_dir`.

// In practice, this enables use a checkpoint for previous experiment for

// initialization.

class TFRecordCheckpointer : public Checkpointer {

public:

explicit TFRecordCheckpointer(

std::string root_dir, std::string group = "",

absl::optional<std::string> fallback_checkpoint_path = absl::nullopt);

// Save a new checkpoint for every table in `tables` in sub directory

// inside `root_dir_`. If the call is successful, the ABSOLUTE path to the

// newly created checkpoint directory is returned.

//

// If `root_path_` does not exist then `Save` attempts to recursively

// create it before proceeding.

//

// If `fallback_checkpoint_path` is provided and iff no checkpoint is found in

// `root_path_`, we attempt to reload the checkpoint with path

// `fallback_checkpoint_path`. This way, we are effectively using the fallback

// checkpoint as a way to initialise the service with a checkpoint generated

// by another experiment.

//

// After a successful save, all but the `keep_latest` most recent checkpoints

// are deleted.

absl::Status Save(std::vector<Table*> tables, int keep_latest,

std::string* path) override;

// Attempts to load a checkpoint stored within `root_dir_`.

absl::Status Load(absl::string_view path, ChunkStore* chunk_store,

std::vector<std::shared_ptr<Table>>* tables) override;

// Finds the most recent checkpoint within `root_dir_` and calls `Load`.

absl::Status LoadLatest(std::vector<std::shared_ptr<Table>>* tables) override;

// Attempts to load the fallback checkpoint. If no fallback_checkpoint_path

// was set or if the no checkpoint found then `NotFoundError` is returned.

absl::Status LoadFallbackCheckpoint(

std::vector<std::shared_ptr<Table>>* tables) override;

// Returns a summary string description.

std::string DebugString() const override;

// TFRecordCheckpointer is neither copyable nor movable.

TFRecordCheckpointer(const TFRecordCheckpointer&) = delete;

TFRecordCheckpointer& operator=(const TFRecordCheckpointer&) = delete;

private:

const std::string root_dir_;

const std::string group_;

absl::optional<std::string> fallback_checkpoint_path_;

};

} // namespace reverb

} // namespace deepmind

#endif // REVERB_CC_PLATFORM_TFRECORD_CHECKPOINTER_H_