async_node.h

Go to the documentation of this file.

/*
 * Copyright 2026 The A11 Authors
 *
 * 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 A11_NODES_ASYNC_NODE_H_
#define A11_NODES_ASYNC_NODE_H_
 
#include <cstddef>
#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <string_view>
#include <utility>
#include <vector>
 
#include <absl/base/thread_annotations.h>
#include <absl/status/status.h>
#include <absl/status/status_macros.h>
#include <absl/status/statusor.h>
#include <absl/time/time.h>
 
#include "a11/concurrency/executor.h"
#include "a11/concurrency/future.h"
#include "a11/data/serializable.h"
#include "a11/data/serialization.h"
#include "a11/data/types.h"
#include "a11/stores/chunk_store.h"
#include "a11/stores/chunk_store_reader.h"
#include "a11/stores/chunk_store_writer.h"
#include "thread/boost_primitives.h"
 
namespace a11::net {
class WireStream;
}  // namespace a11::net
 
namespace a11::nodes {
 
struct FinalizeOptions {
  bool wait = false;
  bool close = true;
  std::optional<std::uint32_t> seq = {};
  std::string_view mimetype = {};
};
 
class AsyncNode : public std::enable_shared_from_this<AsyncNode> {
 public:
  static absl::StatusOr<std::shared_ptr<AsyncNode>> Create(
      std::shared_ptr<stores::ChunkStore> store,
      std::shared_ptr<data::SerializationRegistry> serialization_registry =
          nullptr,
      stores::ChunkStoreReaderOptions reader_options = {},
      stores::ChunkStoreWriterOptions writer_options = {});
 
  ~AsyncNode() = default;
 
  AsyncNode(const AsyncNode&) = delete;
  AsyncNode& operator=(const AsyncNode&) = delete;
 
  absl::StatusOr<std::string> GetId() const;
 
  [[nodiscard]] std::shared_ptr<stores::ChunkStore> GetChunkStore() const;
 
  [[nodiscard]] std::shared_ptr<data::SerializationRegistry>
  serialization_registry() const;
 
  absl::Status SetSerializationRegistry(
      std::shared_ptr<data::SerializationRegistry> registry);
 
  absl::StatusOr<std::shared_ptr<stores::ChunkStoreReader>> reader();
 
  absl::StatusOr<std::shared_ptr<stores::ChunkStoreWriter>> writer();
 
  [[nodiscard]] stores::ChunkStoreReaderOptions GetReaderOptions() const;
 
  absl::Status SetReaderOptions(stores::ChunkStoreReaderOptions options);
 
  absl::Status ResetReader(
      std::optional<stores::ChunkStoreReaderOptions> options = std::nullopt);
 
  [[nodiscard]] stores::ChunkStoreWriterOptions GetWriterOptions() const;
 
  absl::Status SetWriterOptions(stores::ChunkStoreWriterOptions options);
 
  [[nodiscard]] absl::Status GetReaderStatus() const;
 
  [[nodiscard]] absl::Status GetWriterStatus() const;
 
  [[nodiscard]] std::optional<absl::Status> GetWriterAbortStatus() const;
 
  a11::Future<bool> IsWritable();
 
  a11::Future<std::uint32_t> PutChunk(
      data::Chunk chunk, std::optional<std::uint32_t> seq = std::nullopt,
      bool final = false);
 
  a11::Future<std::uint32_t> PutFragment(data::NodeFragment fragment);
 
  template <typename T>
  requires data::HasSerialTypeTag<T> a11::Future<std::uint32_t> PutObject(
      T value, std::string_view mimetype,
      std::optional<std::uint32_t> seq = std::nullopt, bool final = false) {
    std::shared_ptr<data::SerializationRegistry> registry;
    {
      thread::MutexLock lock(&mu_);
      registry = serialization_registry_;
    }
    return PutChunk(
        data::MakeChunkObject<T>(std::move(value), data::SerialTypeTag<T>(),
                                 std::string(mimetype), registry),
        seq, final);
  }
 
  template <typename T>
  a11::Future<std::uint32_t> Put(
      const T& value, std::optional<std::uint32_t> seq = std::nullopt,
      bool final = false, std::string_view mimetype = {}) {
    std::shared_ptr<data::SerializationRegistry> registry;
    {
      thread::MutexLock lock(&mu_);
      registry = serialization_registry_;
    }
    absl::StatusOr<data::Chunk> chunk = registry->ToChunk<T>(value, mimetype);
    if (!chunk.ok()) {
      return a11::FailedFuture<std::uint32_t>(chunk.status());
    }
    return PutChunk(std::move(*chunk), seq, final);
  }
 
  a11::Task Finalize(FinalizeOptions options = {});
 
  a11::Task Finalize(data::Chunk chunk, FinalizeOptions options = {});
 
  template <typename T>
  a11::Task Finalize(const T& value, FinalizeOptions options = {}) {
    std::shared_ptr<data::SerializationRegistry> registry;
    {
      thread::MutexLock lock(&mu_);
      registry = serialization_registry_;
    }
    absl::StatusOr<data::Chunk> chunk =
        registry->ToChunk<T>(value, options.mimetype);
    if (!chunk.ok()) {
      return a11::FailedTask(chunk.status());
    }
    return Finalize(std::move(*chunk), options);
  }
 
  a11::Future<std::vector<std::optional<data::NodeFragment>>> NextFragments(
      size_t limit, absl::Duration timeout = absl::InfiniteDuration());
 
  a11::Future<std::optional<data::NodeFragment>> NextFragment(
      absl::Duration timeout = absl::InfiniteDuration());
 
  a11::Future<std::optional<data::NodeFragment>> NextFragmentRaw(
      absl::Duration timeout = absl::InfiniteDuration());
 
  a11::Future<std::optional<data::Chunk>> NextChunk(
      absl::Duration timeout = absl::InfiniteDuration());
 
  template <typename T>
  a11::Future<std::optional<T>> NextObject(
      absl::Duration timeout = absl::InfiniteDuration(),
      std::vector<std::string> mimetype_patterns = {}) {
    std::shared_ptr<AsyncNode> self = shared_from_this();
    return a11::Submit<std::optional<T>>(
        [self = std::move(self), timeout,
         mimetype_patterns = std::move(
             mimetype_patterns)]() mutable -> absl::StatusOr<std::optional<T>> {
          // Raw, so a chunk carrying a value still is one when it gets here.
          // NextFragment() would have materialised it, which is right for every
          // caller that wants bytes and is exactly what this one does not.
          absl::StatusOr<std::optional<data::NodeFragment>> fragment =
              self->NextFragmentRaw(timeout).Await();
          if (!fragment.ok()) {
            return fragment.status();
          }
          if (!fragment->has_value()) {
            return std::nullopt;
          }
          absl::StatusOr<const data::Chunk*> chunk = (*fragment)->GetChunk();
          if (!chunk.ok()) {
            return chunk.status();
          }
          // The fast path: the producer put a T and this reader wants a T, so
          // nothing was encoded and nothing is decoded. Guarded by the tag, not
          // by type identity -- see a11::data::ChunkObject.
          if ((*chunk)->HasObject()) {
            // `if constexpr`, because NextObject is instantiated for types with
            // no tag at all -- a JSON-native value, a bare string -- and those
            // simply have no fast path to take.
            if constexpr (data::HasSerialTypeTag<T>) {
              if (std::optional<T> taken =
                      data::TryTakeObject<T>(**chunk, data::SerialTypeTag<T>());
                  taken.has_value()) {
                return taken;
              }
            }
            // Carrying something else. Fall through to the bytes, which means
            // producing them now: a reader asking for a different type than the
            // producer wrote is exactly when the wire format earns its keep.
            data::Chunk materialised = **chunk;
            ABSL_RETURN_IF_ERROR(materialised.Materialize());
            std::shared_ptr<data::SerializationRegistry> registry;
            {
              thread::MutexLock lock(&self->mu_);
              registry = self->serialization_registry_;
            }
            absl::StatusOr<T> value =
                registry->FromChunk<T>(materialised, mimetype_patterns);
            if (!value.ok()) {
              return value.status();
            }
            return std::optional<T>(std::move(*value));
          }
          if ((*chunk)->IsNull()) {
            if ((*fragment)->continued) {
              return absl::FailedPreconditionError(
                  "A null stream marker must be final");
            }
            return std::nullopt;
          }
          std::shared_ptr<data::SerializationRegistry> registry;
          {
            thread::MutexLock lock(&self->mu_);
            registry = self->serialization_registry_;
          }
          absl::StatusOr<T> value =
              registry->FromChunk<T>(**chunk, mimetype_patterns);
          if (!value.ok()) {
            return value.status();
          }
          return std::optional<T>(std::move(*value));
        });
  }
 
  a11::Task WaitForBufferToDrain();
 
  a11::Task Close();
 
  a11::Task AbortWithStatus(absl::Status status);
 
  absl::Status AttachStream(std::shared_ptr<net::WireStream> stream);
 
  absl::Status DetachStream(const std::shared_ptr<net::WireStream>& stream);
 
  void CancelReader();
 
  void CancelWriter();
 
  void Cancel();
 
 private:
  AsyncNode(std::shared_ptr<stores::ChunkStore> store,
            std::shared_ptr<data::SerializationRegistry> registry,
            stores::ChunkStoreReaderOptions reader_options,
            stores::ChunkStoreWriterOptions writer_options)
      : store_(std::move(store)),
        serialization_registry_(std::move(registry)),
        reader_options_(reader_options),
        writer_options_(writer_options) {}
 
  const std::shared_ptr<stores::ChunkStore> store_;
  mutable thread::Mutex mu_;
  std::shared_ptr<data::SerializationRegistry> serialization_registry_
      ABSL_GUARDED_BY(mu_);
  stores::ChunkStoreReaderOptions reader_options_ ABSL_GUARDED_BY(mu_);
  stores::ChunkStoreWriterOptions writer_options_ ABSL_GUARDED_BY(mu_);
  std::shared_ptr<stores::ChunkStoreReader> reader_ ABSL_GUARDED_BY(mu_);
  std::shared_ptr<stores::ChunkStoreWriter> writer_ ABSL_GUARDED_BY(mu_);
};
 
}  // namespace a11::nodes
 
#endif  // A11_NODES_ASYNC_NODE_H_