#pragma once

#include "envoy/common/conn_pool.h"

#include "envoy/event/dispatcher.h"

#include "envoy/network/connection.h"

#include "envoy/server/overload/overload_manager.h"

#include "envoy/stats/timespan.h"

#include "envoy/upstream/cluster_manager.h"

#include "source/common/common/debug_recursion_checker.h"

#include "source/common/common/dump_state_utils.h"

#include "source/common/common/linked_object.h"

#include "absl/strings/string_view.h"

#include "fmt/ostream.h"

namespace Envoy {

namespace ConnectionPool {

class ConnPoolImplBase;

// A placeholder struct for whatever data a given connection pool needs to

// successfully attach an upstream connection to a downstream connection.

struct AttachContext {

  // Add a virtual destructor to allow for the dynamic_cast ASSERT in typedContext.

};

// ActiveClient provides a base class for connection pool clients that handles connection timings

// as well as managing the connection timeout.

class ActiveClient : public LinkedObject<ActiveClient>,

                     public Network::ConnectionCallbacks,

                     public Event::DeferredDeletable,

                     protected Logger::Loggable<Logger::Id::pool> {

public:

  ActiveClient(ConnPoolImplBase& parent, uint32_t lifetime_stream_limit,

               uint32_t effective_concurrent_streams, uint32_t concurrent_stream_limit);

  ActiveClient(ConnPoolImplBase& parent, uint32_t lifetime_stream_limit,

               uint32_t concurrent_stream_limit);

  ~ActiveClient() override;

  void releaseResourcesBase();

  // Network::ConnectionCallbacks

  // Called if the connection does not complete within the cluster's connectTimeout()

  void onConnectTimeout();

  // Called if the maximum connection duration is reached.

  void onConnectionDurationTimeout();

  // Returns the concurrent stream limit, accounting for if the total stream limit

  // is less than the concurrent stream limit.

  // Returns the application protocol, or absl::nullopt for TCP.

  virtual absl::optional<Http::Protocol> protocol() const PURE;

  // Initialize upstream read filters. Called when connected.

  virtual void initializeReadFilters() PURE;

  // Closes the underlying connection.

  virtual void

  close(Envoy::Network::ConnectionCloseType type = Envoy::Network::ConnectionCloseType::NoFlush,

        absl::string_view details = "") PURE;

  // Returns the ID of the underlying connection.

  virtual uint64_t id() const PURE;

  // Returns true if this closed with an incomplete stream, for stats tracking/ purposes.

  virtual bool closingWithIncompleteStream() const PURE;

  // Returns the number of active streams on this connection.

  virtual uint32_t numActiveStreams() const PURE;

  // Return true if it is ready to dispatch the next stream.

  enum class State {

    Connecting,        // Connection is not yet established.

    ReadyForEarlyData, // Any additional early data stream can be immediately dispatched to this

                       // connection.

    Ready,             // Additional streams may be immediately dispatched to this connection.

    Busy,              // Connection is at its concurrent stream limit.

    Draining,          // No more streams can be dispatched to this connection, and it will be

                       // closed when all streams complete.

    Closed             // Connection is closed and object is queued for destruction.

  };

    // If the client is transitioning to draining, update the remaining

    // streams and pool and cluster capacity.

  // Sets the remaining streams to 0, and updates pool and cluster capacity.

  virtual void drain();

  ConnPoolImplBase& parent_;

  // The count of remaining streams allowed for this connection.

  // This will start out as the total number of streams per connection if capped

  // by configuration, or it will be set to std::numeric_limits<uint32_t>::max() to be

  // (functionally) unlimited.

  // TODO: this could be moved to an optional to make it actually unlimited.

  uint32_t remaining_streams_;

  // The will start out as the upper limit of max concurrent streams for this connection

  // if capped by configuration, or it will be set to std::numeric_limits<uint32_t>::max()

  // to be (functionally) unlimited.

  uint32_t configured_stream_limit_;

  // The max concurrent stream for this connection, it's initialized by `configured_stream_limit_`

  // and can be adjusted by SETTINGS frame, but the max value of it can't exceed

  // `configured_stream_limit_`.

  uint32_t concurrent_stream_limit_;

  Upstream::HostDescriptionConstSharedPtr real_host_description_;

  Stats::TimespanPtr conn_connect_ms_;

  Stats::TimespanPtr conn_length_;

  Event::TimerPtr connect_timer_;

  Event::TimerPtr connection_duration_timer_;

  bool resources_released_{false};

  bool timed_out_{false};

  // TODO(danzh) remove this once http codec exposes the handshake state for h3.

  bool has_handshake_completed_{false};

protected:

  // HTTP/3 subclass should override this.

private:

  State state_{State::Connecting};

};

// PendingStream is the base class tracking streams for which a connection has been created but not

// yet established.

class PendingStream : public LinkedObject<PendingStream>, public ConnectionPool::Cancellable {

public:

  PendingStream(ConnPoolImplBase& parent, bool can_send_early_data);

  ~PendingStream() override;

  // ConnectionPool::Cancellable

  void cancel(Envoy::ConnectionPool::CancelPolicy policy) override;

  // The context here returns a pointer to whatever context is provided with newStream(),

  // which will be passed back to the parent in onPoolReady or onPoolFailure.

  virtual AttachContext& context() PURE;

  ConnPoolImplBase& parent_;

  // The request can be sent as early data.

  bool can_send_early_data_;

};

using PendingStreamPtr = std::unique_ptr<PendingStream>;

using ActiveClientPtr = std::unique_ptr<ActiveClient>;

// Base class that handles stream queueing logic shared between connection pool implementations.

class ConnPoolImplBase : protected Logger::Loggable<Logger::Id::pool> {

public:

  ConnPoolImplBase(Upstream::HostConstSharedPtr host, Upstream::ResourcePriority priority,

                   Event::Dispatcher& dispatcher,

                   const Network::ConnectionSocket::OptionsSharedPtr& options,

                   const Network::TransportSocketOptionsConstSharedPtr& transport_socket_options,

                   Upstream::ClusterConnectivityState& state,

                   Server::OverloadManager& overload_manager);

  virtual ~ConnPoolImplBase();

  void deleteIsPendingImpl();

  // By default, the connection pool will track connected and connecting stream

  // capacity as streams are created and destroyed. QUIC does custom stream

  // accounting so will override this to false.

  // A helper function to get the specific context type from the base class context.

  // Determines if prefetching is warranted based on the number of streams in

  // use, pending streams, anticipated and/or currently unused capacity, and

  // preconnect configuration.

  //

  // If anticipate_incoming_stream is true this assumes a call to newStream is

  // pending, which is true for global preconnect.

  static bool shouldConnect(size_t pending_streams, size_t active_streams,

                            int64_t connecting_and_connected_capacity, float preconnect_ratio,

                            bool anticipate_incoming_stream = false);

  // Envoy::ConnectionPool::Instance implementation helpers

  void addIdleCallbackImpl(Instance::IdleCb cb);

  // Returns true if the pool is idle.

  bool isIdleImpl() const;

  void drainConnectionsImpl(DrainBehavior drain_behavior);

  // Called if this pool is likely to be picked soon, to determine if it's worth preconnecting.

  bool maybePreconnectImpl(float global_preconnect_ratio);

  // Closes and destroys all connections. This must be called in the destructor of

  // derived classes because the derived ActiveClient will downcast parent_ to a more

  // specific type of ConnPoolImplBase, but if the more specific part is already destructed

  // (due to bottom-up destructor ordering in c++) that access will be invalid.

  void destructAllConnections();

  // Returns a new instance of ActiveClient.

  virtual ActiveClientPtr instantiateActiveClient() PURE;

  // Gets a pointer to the list that currently owns this client.

  std::list<ActiveClientPtr>& owningList(ActiveClient::State state);

  // Removes the PendingStream from the list of streams. Called when the PendingStream is

  // cancelled, e.g. when the stream is reset before a connection has been established.

  void onPendingStreamCancel(PendingStream& stream, Envoy::ConnectionPool::CancelPolicy policy);

  // Fails all pending streams, calling onPoolFailure on the associated callbacks.

  void purgePendingStreams(const Upstream::HostDescriptionConstSharedPtr& host_description,

                           absl::string_view failure_reason,

                           ConnectionPool::PoolFailureReason pool_failure_reason);

  // Closes any idle connections as this pool is drained.

  void closeIdleConnectionsForDrainingPool();

  // Changes the state_ of an ActiveClient and moves to the appropriate list.

  void transitionActiveClientState(ActiveClient& client, ActiveClient::State new_state);

  void onConnectionEvent(ActiveClient& client, absl::string_view failure_reason,

                         Network::ConnectionEvent event);

  // Check if the pool has gone idle and invoke idle notification callbacks.

  void checkForIdleAndNotify();

  // See if the pool has gone idle. If we're draining, this will also close idle connections.

  void checkForIdleAndCloseIdleConnsIfDraining();

  void scheduleOnUpstreamReady();

  ConnectionPool::Cancellable* newStreamImpl(AttachContext& context, bool can_send_early_data);

  virtual ConnectionPool::Cancellable* newPendingStream(AttachContext& context,

                                                        bool can_send_early_data) PURE;

  virtual void attachStreamToClient(Envoy::ConnectionPool::ActiveClient& client,

                                    AttachContext& context);

  virtual void onPoolFailure(const Upstream::HostDescriptionConstSharedPtr& host_description,

                             absl::string_view failure_reason,

                             ConnectionPool::PoolFailureReason pool_failure_reason,

                             AttachContext& context) PURE;

  virtual void onPoolReady(ActiveClient& client, AttachContext& context) PURE;

  // Called by derived classes any time a stream is completed or destroyed for any reason.

  void onStreamClosed(Envoy::ConnectionPool::ActiveClient& client, bool delay_attaching_stream);

  // Helper for use as the 2nd argument to ASSERT.

  std::string dumpState() const;

  void decrConnectingAndConnectedStreamCapacity(uint32_t delta, ActiveClient& client);

  void incrConnectingAndConnectedStreamCapacity(uint32_t delta, ActiveClient& client);

  // Called when an upstream is ready to serve pending streams.

  void onUpstreamReady();

  // Called when an upstream is ready to serve early data streams.

  void onUpstreamReadyForEarlyData(ActiveClient& client);

protected:

  enum class ConnectionResult {

    FailedToCreateConnection,

    CreatedNewConnection,

    ShouldNotConnect,

    NoConnectionRateLimited,

    CreatedButRateLimited,

    LoadShed,

  };

  // Creates up to 3 connections, based on the preconnect ratio.

  // Returns the ConnectionResult of the last attempt.

  ConnectionResult tryCreateNewConnections();

  // Creates a new connection if there is sufficient demand, it is allowed by resourceManager, or

  // to avoid starving this pool.

  // Demand is determined either by perUpstreamPreconnectRatio() or global_preconnect_ratio

  // if this is called by maybePreconnect()

  ConnectionResult tryCreateNewConnection(float global_preconnect_ratio = 0);

  // A helper function which determines if a canceled pending connection should

  // be closed as excess or not.

  bool connectingConnectionIsExcess(const ActiveClient& client) const;

  // A helper function which determines if a new incoming stream should trigger

  // connection preconnect.

  bool shouldCreateNewConnection(float global_preconnect_ratio) const;

  float perUpstreamPreconnectRatio() const;

  ConnectionPool::Cancellable*

  const Upstream::HostConstSharedPtr host_;

  const Upstream::ResourcePriority priority_;

  Event::Dispatcher& dispatcher_;

  const Network::ConnectionSocket::OptionsSharedPtr socket_options_;

  const Network::TransportSocketOptionsConstSharedPtr transport_socket_options_;

  // True if the max requests circuit breakers apply.

  // This will be false for the TCP pool, true otherwise.

  std::list<Instance::IdleCb> idle_callbacks_;

  // When calling purgePendingStreams, this list will be used to hold the streams we are about

  // to purge. We need this if one cancelled streams cancels a different pending stream

  std::list<PendingStreamPtr> pending_streams_to_purge_;

  // Clients that are ready to handle additional streams.

  // All entries are in state Ready.

  std::list<ActiveClientPtr> ready_clients_;

  // Clients that are not ready to handle additional streams due to being Busy or Draining.

  std::list<ActiveClientPtr> busy_clients_;

  // Clients that are not ready to handle additional streams because they are Connecting.

  std::list<ActiveClientPtr> connecting_clients_;

  // Clients that are ready to handle additional early data streams because they have 0-RTT

  // credentials.

  std::list<ActiveClientPtr> early_data_clients_;

  // The number of streams that can be immediately dispatched

  // if all Connecting connections become connected.

  uint32_t connecting_stream_capacity_{0};

private:

  // Drain all the clients in the given list.

  // Prerequisite: the given clients shouldn't be idle.

  void drainClients(std::list<ActiveClientPtr>& clients);

  void assertCapacityCountsAreCorrect();

  Upstream::ClusterConnectivityState& cluster_connectivity_state_;

  std::list<PendingStreamPtr> pending_streams_;

  // The number of streams that can be immediately dispatched from the current

  // `ready_clients_` plus `connecting_stream_capacity_`.

  int64_t connecting_and_connected_stream_capacity_{0};

  // The number of streams currently attached to clients.

  uint32_t num_active_streams_{0};

  // Whether the connection pool is currently in the process of closing

  // all connections so that it can be gracefully deleted.

  bool is_draining_for_deletion_{false};

  // True iff this object is in the deferred delete list.

  bool deferred_deleting_{false};

  Event::SchedulableCallbackPtr upstream_ready_cb_;

  Common::DebugRecursionChecker recursion_checker_;

  Server::LoadShedPoint* create_new_connection_load_shed_{nullptr};

};

} // namespace ConnectionPool

} // namespace Envoy

// NOLINT(namespace-envoy)

namespace fmt {

template <> struct formatter<Envoy::ConnectionPool::ConnPoolImplBase> : ostream_formatter {};

} // namespace fmt