chunk_store_reader.h

// Copyright 2025 Google LLC
//
// 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 ACTIONENGINE_STORES_CHUNK_STORE_READER_H_
#define ACTIONENGINE_STORES_CHUNK_STORE_READER_H_
 
#include <algorithm>
#include <cstddef>
#include <memory>
#include <optional>
#include <utility>
#include <vector>
 
#include <absl/base/nullability.h>
#include <absl/base/thread_annotations.h>
#include <absl/log/check.h>
#include <absl/status/status.h>
#include <absl/status/statusor.h>
#include <absl/time/time.h>
 
#include "actionengine/concurrency/concurrency.h"
#include "actionengine/data/serialization.h"
#include "actionengine/data/types.h"
#include "actionengine/stores/chunk_store.h"
#include "actionengine/util/status_macros.h"
 
namespace act {

struct ChunkStoreReaderOptions {
  bool ordered = true;

  bool remove_chunks = true;

  size_t n_chunks_to_buffer = 32;

  absl::Duration timeout = absl::InfiniteDuration();
 
  size_t start_seq_or_offset = 0;
};

class ChunkStoreReader {
 public:
  explicit ChunkStoreReader(
      ChunkStore* absl_nonnull chunk_store,
      ChunkStoreReaderOptions options = ChunkStoreReaderOptions());
 
  // This class is not copyable or movable.
  ChunkStoreReader(const ChunkStoreReader&) = delete;
  ChunkStoreReader& operator=(const ChunkStoreReader&) = delete;
 
  ~ChunkStoreReader();

  void Cancel() const;

  void SetOptions(const ChunkStoreReaderOptions& options);

  [[nodiscard]]
  const ChunkStoreReaderOptions& GetOptions() const {
    return options_;
  }

  absl::StatusOr<std::optional<Chunk>> Next(
      std::optional<absl::Duration> timeout = std::nullopt);
 
  absl::StatusOr<std::optional<NodeFragment>> NextFragment(
      std::optional<absl::Duration> timeout = std::nullopt) {
    act::MutexLock lock(&mu_);
    absl::StatusOr<std::optional<std::pair<int, Chunk>>> seq_and_chunk =
        GetNextSeqAndChunkFromBuffer(timeout.value_or(options_.timeout));
    RETURN_IF_ERROR(seq_and_chunk.status());
    ASSIGN_OR_RETURN(const int64_t final_seq, chunk_store_->GetFinalSeq());
    if (!seq_and_chunk->has_value()) {
      return std::nullopt;
    }
    if (seq_and_chunk->value().second.IsNull()) {
      // If the chunk is null, it means that the stream has ended.
      // TODO: this logic is not ideal, as it does not allow to distinguish
      //   between an empty chunk and the end of the stream. We should rethink it.
      return std::nullopt;
    }
    auto& chunk = seq_and_chunk->value().second;
    const int seq = seq_and_chunk->value().first;
    return NodeFragment{std::string(chunk_store_->GetId()), std::move(chunk),
                        seq, final_seq == -1 || seq != final_seq};
  }

  template <typename T>
  absl::StatusOr<std::optional<T>> Next(
      std::optional<absl::Duration> timeout = std::nullopt) {
    ASSIGN_OR_RETURN(std::optional<Chunk> chunk, Next(timeout));
    if (!chunk) {
      return std::nullopt;
    }
    if (chunk->metadata && chunk->metadata->mimetype == "__status__") {
      absl::StatusOr<absl::Status> status = ConvertTo<absl::Status>(*chunk);
      if (!status.ok()) {
        return status.status();
      }
      // Nested status
      if (!status->ok()) {
        return *status;
      }
      return absl::FailedPreconditionError(
          "Next<T>() cannot return an OK status chunk if T is not "
          "absl::Status. ");
    }
    ASSIGN_OR_RETURN(T result, FromChunkAs<T>(*std::move(chunk)));
    return std::optional{result};
  }
 
  // Definitions follow in the header for some well-known types. If the next
  // chunk cannot be read, this operator will crash. Prefer to use the Next()
  // method instead, which returns an absl::StatusOr<std::optional<T>>,
  // unless you are sure that the next chunk will always be available.
  template <typename T>
  friend ChunkStoreReader& operator>>(ChunkStoreReader& reader, T& value);

  absl::Status GetStatus() const;
 
 private:
  friend class Action;
 
  void EnsurePrefetchIsRunningOrHasCompleted()
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
 
  absl::StatusOr<std::optional<Chunk>> GetNextChunkFromBuffer(
      absl::Duration timeout) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
  absl::StatusOr<std::optional<std::pair<int, Chunk>>>
  GetNextSeqAndChunkFromBuffer(absl::Duration timeout)
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
 
  absl::StatusOr<std::optional<std::pair<int, Chunk>>>
  GetNextUnorderedSeqAndChunkFromStore() const
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
 
  absl::Status RunPrefetchLoop() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
 
  ChunkStore* const absl_nonnull chunk_store_;
  ChunkStoreReaderOptions options_;
 
  std::unique_ptr<thread::Fiber> fiber_;
  std::unique_ptr<thread::Channel<std::optional<std::pair<int, Chunk>>>>
      buffer_;
  int total_chunks_read_ = 0;
 
  absl::Status status_;
  mutable act::Mutex mu_;
};
 
// These specializations simply take the mutex and call the internal methods
// GetNextSeqAndChunkFromBuffer and GetNextChunkFromBuffer.
template <>
absl::StatusOr<std::optional<std::pair<int, Chunk>>> ChunkStoreReader::Next(
    std::optional<absl::Duration> timeout);
 
template <>
absl::StatusOr<std::optional<absl::Status>> ChunkStoreReader::Next(
    std::optional<absl::Duration> timeout);
 
template <typename T>
ChunkStoreReader& operator>>(ChunkStoreReader& reader,
                             std::optional<T>& value) {
  absl::StatusOr<std::optional<T>> next_value = reader.Next<T>();
  CHECK_OK(next_value.status())
      << "Failed to read next value: " << next_value.status();
  value = *std::move(next_value);
  return reader;
}
 
template <typename T>
ChunkStoreReader& operator>>(ChunkStoreReader& reader, std::vector<T>& value) {
  while (true) {
    absl::StatusOr<std::optional<T>> status_or_element = reader.Next<T>();
    CHECK_OK(status_or_element.status())
        << "Failed to read next element: " << status_or_element.status();
 
    std::optional<T>& element = *status_or_element;
    if (!element) {
      break;
    }
 
    value.push_back(*std::move(element));
  }
  return reader;
}
 
template <typename T>
ChunkStoreReader* absl_nonnull operator>>(ChunkStoreReader* absl_nonnull reader,
                                          T& value) {
  *reader >> value;
  return reader;
}
 
template <typename T>
std::unique_ptr<ChunkStoreReader>& operator>>(
    std::unique_ptr<ChunkStoreReader>& reader, T& value) {
  *reader >> value;
  return reader;
}
 
template <typename T>
std::shared_ptr<ChunkStoreReader>& operator>>(
    std::shared_ptr<ChunkStoreReader>& reader, T& value) {
  *reader >> value;
  return reader;
}
 
}  // namespace act
 
#endif  // ACTIONENGINE_STORES_CHUNK_STORE_READER_H_