stream.h

#pragma once

#include <array>
#include <cstdint>
#include <memory>
#include <queue>
#include <functional>
#include <string_view>
#include <type_traits>
#include <variant>
#include <vector>

#include <uvw/async.h>

namespace quic {

class Connection;

using bstring_view = std::basic_string_view<std::byte>;

// Shortcut for a const-preserving `reinterpret_cast`ing c.data() from a std::byte to a uint8_t
// pointer, because we need it all over the place in the ngtcp2 API and I'd rather deal with
// std::byte's out here for type safety.
template <typename Container, typename = std::enable_if_t<sizeof(typename std::remove_reference_t<Container>::value_type) == sizeof(uint8_t)>>
inline auto* u8data(Container&& c) {
    using u8_sameconst_t = std::conditional_t<std::is_const_v<std::remove_pointer_t<decltype(c.data())>>,
          const uint8_t, uint8_t>;
    return reinterpret_cast<u8_sameconst_t*>(c.data());
}

// Type-safe wrapper around a int64_t stream id.  Default construction is ngtcp2's special
// "no-stream" id.
struct StreamID {
    int64_t id{-1};
    bool operator==(const StreamID &s) const { return s.id == id; }
    bool operator!=(const StreamID &s) const { return s.id != id; }
    bool operator<(const StreamID &s) const { return s.id < id; }
    bool operator<=(const StreamID &s) const { return s.id <= id; }
    bool operator>(const StreamID &s) const { return s.id > id; }
    bool operator>=(const StreamID &s) const { return s.id >= id; }
};

// Application error code we close with if the data handle throws
constexpr uint64_t STREAM_EXCEPTION_ERROR_CODE = (1ULL << 62) - 2;

std::ostream& operator<<(std::ostream& o, const StreamID& s);
} // namespace quic

namespace std {
template <> struct hash<quic::StreamID> {
    size_t operator()(const quic::StreamID& s) const {
        return std::hash<decltype(s.id)>{}(s.id);
    }
};
}

namespace quic {

// Class for an established stream (a single connection has multiple streams): we have a fixed-sized
// ring buffer for holding outgoing data, and a callback to invoke on received data.  To construct a
// Stream call `conn.open_stream()`.
class Stream : public std::enable_shared_from_this<Stream> {
public:

    // Returns the StreamID of this stream
    const StreamID& id() const { return stream_id; }

    // Sets the size of the outgoing data buffer.  This may *only* be used if the buffer is
    // currently entirely empty; otherwise a runtime_error is thrown.  The minimum buffer size is
    // 2048, the default is 64kiB.  A value of 0 puts the Stream into user-provided buffer mode
    // where only the version of `append` taking ownership of a char* is permitted.
    void set_buffer_size(size_t size);

    // Returns the size of the buffer (including both pending and free space).  If using
    // user-provided buffer mode then this is the sum of all held buffers.
    size_t buffer_size() const;

    // Returns the number of free bytes available in the outgoing stream data buffer. Always 0 in
    // user-provided buffer mode.
    size_t available() const { return is_closing || buffer.empty() ? 0 : buffer.size() - size; }

    // Returns the number of bytes currently referenced in the buffer (i.e. pending or
    // sent-but-unacknowledged).
    size_t used() const { return size; }

    // Returns the number of bytes of the buffer that have been sent but not yet acknowledged and
    // thus are still required.
    size_t unacked() const { return unacked_size; }

    // Returns the number of bytes of the buffer that have not yet been sent.
    size_t unsent() const { return used() - unacked(); }

    // Try to append all of the given bytes to the outgoing stream data buffer.  Returns true if
    // successful, false (without appending anything) if there is insufficient space.  If you want
    // to append as much as possible then use `append_any` instead.
    bool append(bstring_view data);
    bool append(std::string_view data) {
        return append(bstring_view{reinterpret_cast<const std::byte*>(data.data()), data.size()});
    }

    // Append bytes to the outgoing stream data buffer, allowing partial consumption of data if the
    // entire provided data cannot be appended.  Returns the number of appended bytes (which will be
    // less than the total provided if the provided data is larger than `available()`).  If you want
    // an all-or-nothing append then use `append` instead.
    size_t append_any(bstring_view data);
    size_t append_any(std::string_view data) {
        return append_any(bstring_view{reinterpret_cast<const std::byte*>(data.data()), data.size()});
    }

    // Takes ownership of the given buffer pointer, queuing it to be sent after any existing buffers
    // and freed once fully acked.  You *must* have called `set_buffer_size(0)` (or set the
    // endpoints default_stream_buffer_size to 0) in order to use this.
    void append_buffer(const std::byte* buf, size_t length);

    // Starting closing the stream and prevent any more outgoing data from being appended.  If
    // `error_code` is provided then we close immediately with the given code; if std::nullopt (the
    // default) we close gracefully by sending a FIN bit.
    void close(std::optional<uint64_t> error_code = std::nullopt);

    // Returns true if this Stream is closing (or already closed).
    bool closing() const { return is_closing; }

    // Callback invoked when data is received
    using data_callback_t = std::function<void(Stream&, bstring_view)>;

    // Callback invoked when the stream is closed
    using close_callback_t = std::function<void(Stream&, std::optional<uint64_t> error_code)>;

    // Callback invoked when free stream buffer space becomes available.  Should return true if the
    // callback is finished and can be discarded, false if the callback is still needed.  If
    // returning false then it *must* have filled the stream's outgoing buffer (this is asserted in
    // a debug build).
    using unblocked_callback_t = std::function<bool(Stream&)>;

    // Callback to invoke when we receive some incoming data; there's no particular guarantee on the
    // size of the data, just that this will always be called in sequential order.
    data_callback_t data_callback;

    // Callback to invoke when the connection has closed.  If the close was an abrupt stream close
    // initiated by the remote then `error_code` will be set to whatever code the remote side
    // provided; for graceful closing or locally initiated closing the error code will be null.
    close_callback_t close_callback;

    // Queues a callback to be invoked when space becomes available for writing in the buffer.  The
    // callback should true if it completed, false if it still needs more buffer space.  If multiple
    // callbacks are queued they are invoked in order, space permitting.  The stored std::function
    // will not be moved or copied after being invoked (i.e. if invoked multiple times it will
    // always be invoked on the same instance).
    //
    // Available callbacks should only be used when the buffer is full, typically immediately after
    // an `append_any` call that returns less than the full write.  Similarly a false return from an
    // unblock function (which keeps the callback alive) should satisfy the same condition.
    //
    // In user-provided buffer mode the callback will be invoked after any data has been acked: it
    // is up to the caller to look at used()/buffer_size()/etc. to decide what to do.  As described
    // above, return true to remove this callback, false to keep it and try again after the next
    // ack.
    void when_available(unblocked_callback_t unblocked_cb);

    // Calls io_ready() on the stream's connection to scheduling sending outbound data
    void io_ready();

    // Schedules processing of the "when_available" callbacks
    void available_ready();

    // Lets you stash some arbitrary data in a shared_ptr; this is not used internally.
    void data(std::shared_ptr<void> data);

    // Variation of data() that holds the pointer in a weak_ptr instead of a shared_ptr.
    void weak_data(std::weak_ptr<void> data);

    // Retrieves the stashed data, with a static_cast to the desired type.  This is used for
    // retrieval of both shared or weak data types (if held as a weak_ptr it is lock()ed first).
    template <typename T>
    std::shared_ptr<T> data() const {
        return std::static_pointer_cast<T>(
            std::holds_alternative<std::shared_ptr<void>>(user_data)
                ? std::get<std::shared_ptr<void>>(user_data)
                : std::get<std::weak_ptr<void>>(user_data).lock());
    }

private:
    friend class Connection;

    Stream(Connection& conn, data_callback_t data_cb, close_callback_t close_cb, size_t buffer_size, StreamID id = {-1});
    Stream(Connection& conn, StreamID id, size_t buffer_size);

    // Non-copyable, non-movable; we manage it via a unique_ptr held by its Connection
    Stream(const Stream&) = delete;
    const Stream& operator=(const Stream&) = delete;
    Stream(Stream&&) = delete;
    Stream& operator=(Stream&&) = delete;

    Connection& conn;

    // Callback(s) to invoke once we have the requested amount of space available in the buffer.
    std::queue<unblocked_callback_t> unblocked_callbacks;
    void handle_unblocked(); // Processes the above if space is available

    // Called to advance the number of acknowledged bytes (freeing up that space in the buffer for
    // appending data).
    void acknowledge(size_t bytes);

    // Returns a view into unwritten stream data.  This returns a vector of string_views of the data
    // to write, in order.  After writing any of the provided data you must call `wrote()` to signal
    // how much of the given data was consumed (to advance the next pending() call).
    std::vector<bstring_view> pending();

    // Called to signal that bytes have been written and should now be considered sent (but still
    // unacknowledged), thereby advancing the initial data position returned by the next `pending()`
    // call.  Should typically be called after `pending()` to signal how much of the pending data
    // was actually used.
    void wrote(size_t bytes);

    // ngtcp2 stream_id, assigned during stream creation
    StreamID stream_id{-1};

    // ring buffer of outgoing stream data that has not yet been acknowledged.  This cannot be
    // resized once used as ngtcp2 will have pointers into the data.  If this is empty then we are
    // in user-provided buffer mode.
    std::vector<std::byte> buffer{65536};

    // user-provided buffers; only used when `buffer` is empty (via a `set_buffer_size(0)` or a 0
    // size given in the constructor).
    std::deque<std::pair<std::unique_ptr<const std::byte[]>, size_t>> user_buffers;

    // Offset of the first used byte in the circular buffer, will always be in [0, buffer.size()).
    // For user-provided buffers this is the starting offset in the currently sending user-provided
    // buffer.
    size_t start{0};

    // Number of sent-but-unacked packets in the buffer (i.e. [start, start+unacked_size) are sent but
    // not yet acked).
    size_t unacked_size{0};

    // Number of used bytes in the buffer; thus start+size is the next write location and
    // [start+unacked_size, start+size) is the range of not-yet-sent bytes.  (Note that this
    // description is ignoring the circularity of the buffer).
    size_t size{0};

    bool is_new{true};
    bool is_closing{false};
    bool sent_fin{false};
    bool is_shutdown{false};

    // Async trigger we use to schedule when_available callbacks (so that we can make them happen in
    // batches rather than after each and every packet ack).
    std::shared_ptr<uvw::AsyncHandle> avail_trigger;

    std::variant<std::shared_ptr<void>, std::weak_ptr<void>> user_data;
};

} // namespace quic