// Copyright (c) 2010 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
//
// A ClientSocketPoolBase is used to restrict the number of sockets open at
// a time.  It also maintains a list of idle persistent sockets for reuse.
// Subclasses of ClientSocketPool should compose ClientSocketPoolBase to handle
// the core logic of (1) restricting the number of active (connected or
// connecting) sockets per "group" (generally speaking, the hostname), (2)
// maintaining a per-group list of idle, persistent sockets for reuse, and (3)
// limiting the total number of active sockets in the system.
//
// ClientSocketPoolBase abstracts socket connection details behind ConnectJob,
// ConnectJobFactory, and SocketParams.  When a socket "slot" becomes available,
// the ClientSocketPoolBase will ask the ConnectJobFactory to create a
// ConnectJob with a SocketParams.  Subclasses of ClientSocketPool should
// implement their socket specific connection by subclassing ConnectJob and
// implementing ConnectJob::ConnectInternal().  They can control the parameters
// passed to each new ConnectJob instance via their ConnectJobFactory subclass
// and templated SocketParams parameter.
//
#ifndef NET_SOCKET_CLIENT_SOCKET_POOL_BASE_H_
#define NET_SOCKET_CLIENT_SOCKET_POOL_BASE_H_
#pragma once

#include <deque>
#include <map>
#include <set>
#include <string>

#include "base/basictypes.h"
#include "base/compiler_specific.h"
#include "base/ref_counted.h"
#include "base/scoped_ptr.h"
#include "base/task.h"
#include "base/time.h"
#include "base/timer.h"
#include "net/base/address_list.h"
#include "net/base/completion_callback.h"
#include "net/base/load_states.h"
#include "net/base/net_errors.h"
#include "net/base/net_log.h"
#include "net/base/network_change_notifier.h"
#include "net/base/request_priority.h"
#include "net/socket/client_socket.h"
#include "net/socket/client_socket_pool.h"

namespace net {

class ClientSocketHandle;

// ConnectJob provides an abstract interface for "connecting" a socket.
// The connection may involve host resolution, tcp connection, ssl connection,
// etc.
class ConnectJob {
 public:
  class Delegate {
   public:
    Delegate() {}
    virtual ~Delegate() {}

    // Alerts the delegate that the connection completed.
    virtual void OnConnectJobComplete(int result, ConnectJob* job) = 0;

   private:
    DISALLOW_COPY_AND_ASSIGN(Delegate);
  };

  // A |timeout_duration| of 0 corresponds to no timeout.
  ConnectJob(const std::string& group_name,
             base::TimeDelta timeout_duration,
             Delegate* delegate,
             const BoundNetLog& net_log);
  virtual ~ConnectJob();

  // Accessors
  const std::string& group_name() const { return group_name_; }
  const BoundNetLog& net_log() { return net_log_; }

  // Releases |socket_| to the client.  On connection error, this should return
  // NULL.
  ClientSocket* ReleaseSocket() { return socket_.release(); }

  // Begins connecting the socket.  Returns OK on success, ERR_IO_PENDING if it
  // cannot complete synchronously without blocking, or another net error code
  // on error.  In asynchronous completion, the ConnectJob will notify
  // |delegate_| via OnConnectJobComplete.  In both asynchronous and synchronous
  // completion, ReleaseSocket() can be called to acquire the connected socket
  // if it succeeded.
  int Connect();

  virtual LoadState GetLoadState() const = 0;

  // If Connect returns an error (or OnConnectJobComplete reports an error
  // result) this method will be called, allowing the pool to add
  // additional error state to the ClientSocketHandle (post late-binding).
  virtual void GetAdditionalErrorState(ClientSocketHandle* handle) {}

 protected:
  void set_socket(ClientSocket* socket);
  ClientSocket* socket() { return socket_.get(); }
  void NotifyDelegateOfCompletion(int rv);
  void ResetTimer(base::TimeDelta remainingTime);

 private:
  virtual int ConnectInternal() = 0;

  void LogConnectStart();
  void LogConnectCompletion(int net_error);

  // Alerts the delegate that the ConnectJob has timed out.
  void OnTimeout();

  const std::string group_name_;
  const base::TimeDelta timeout_duration_;
  // Timer to abort jobs that take too long.
  base::OneShotTimer<ConnectJob> timer_;
  Delegate* delegate_;
  scoped_ptr<ClientSocket> socket_;
  BoundNetLog net_log_;
  // A ConnectJob is idle until Connect() has been called.
  bool idle_;

  DISALLOW_COPY_AND_ASSIGN(ConnectJob);
};

namespace internal {

// ClientSocketPoolBaseHelper is an internal class that implements almost all
// the functionality from ClientSocketPoolBase without using templates.
// ClientSocketPoolBase adds templated definitions built on top of
// ClientSocketPoolBaseHelper.  This class is not for external use, please use
// ClientSocketPoolBase instead.
class ClientSocketPoolBaseHelper
    : public base::RefCounted<ClientSocketPoolBaseHelper>,
      public ConnectJob::Delegate,
      public NetworkChangeNotifier::Observer {
 public:
  class Request {
   public:
    Request(ClientSocketHandle* handle,
            CompletionCallback* callback,
            RequestPriority priority,
            const BoundNetLog& net_log);

    virtual ~Request();

    ClientSocketHandle* handle() const { return handle_; }
    CompletionCallback* callback() const { return callback_; }
    RequestPriority priority() const { return priority_; }
    const BoundNetLog& net_log() const { return net_log_; }

   private:
    ClientSocketHandle* const handle_;
    CompletionCallback* const callback_;
    const RequestPriority priority_;
    BoundNetLog net_log_;

    DISALLOW_COPY_AND_ASSIGN(Request);
  };

  class ConnectJobFactory {
   public:
    ConnectJobFactory() {}
    virtual ~ConnectJobFactory() {}

    virtual ConnectJob* NewConnectJob(
        const std::string& group_name,
        const Request& request,
        ConnectJob::Delegate* delegate) const = 0;

    virtual base::TimeDelta ConnectionTimeout() const = 0;

   private:
    DISALLOW_COPY_AND_ASSIGN(ConnectJobFactory);
  };

  ClientSocketPoolBaseHelper(
      int max_sockets,
      int max_sockets_per_group,
      base::TimeDelta unused_idle_socket_timeout,
      base::TimeDelta used_idle_socket_timeout,
      ConnectJobFactory* connect_job_factory);

  // See ClientSocketPool::RequestSocket for documentation on this function.
  // ClientSocketPoolBaseHelper takes ownership of |request|, which must be
  // heap allocated.
  int RequestSocket(const std::string& group_name, const Request* request);

  // See ClientSocketPool::CancelRequest for documentation on this function.
  void CancelRequest(const std::string& group_name,
                     ClientSocketHandle* handle);

  // See ClientSocketPool::ReleaseSocket for documentation on this function.
  void ReleaseSocket(const std::string& group_name,
                     ClientSocket* socket,
                     int id);

  // See ClientSocketPool::Flush for documentation on this function.
  void Flush();

  // See ClientSocketPool::CloseIdleSockets for documentation on this function.
  void CloseIdleSockets();

  // See ClientSocketPool::IdleSocketCount() for documentation on this function.
  int idle_socket_count() const {
    return idle_socket_count_;
  }

  // See ClientSocketPool::IdleSocketCountInGroup() for documentation on this
  // function.
  int IdleSocketCountInGroup(const std::string& group_name) const;

  // See ClientSocketPool::GetLoadState() for documentation on this function.
  LoadState GetLoadState(const std::string& group_name,
                         const ClientSocketHandle* handle) const;

  int ConnectRetryIntervalMs() const {
    // TODO(mbelshe): Make this tuned dynamically based on measured RTT.
    //                For now, just use the max retry interval.
    return ClientSocketPool::kMaxConnectRetryIntervalMs;
  }

  // ConnectJob::Delegate methods:
  virtual void OnConnectJobComplete(int result, ConnectJob* job);

  // NetworkChangeNotifier::Observer methods:
  virtual void OnIPAddressChanged();

  int NumConnectJobsInGroup(const std::string& group_name) const {
    return group_map_.find(group_name)->second.jobs.size();
  }

  // Closes all idle sockets if |force| is true.  Else, only closes idle
  // sockets that timed out or can't be reused.  Made public for testing.
  void CleanupIdleSockets(bool force);

  base::TimeDelta ConnectionTimeout() const {
    return connect_job_factory_->ConnectionTimeout();
  }

  void EnableBackupJobs() { backup_jobs_enabled_ = true; }

 private:
  friend class base::RefCounted<ClientSocketPoolBaseHelper>;

  // Entry for a persistent socket which became idle at time |start_time|.
  struct IdleSocket {
    IdleSocket() : socket(NULL), used(false) {}
    ClientSocket* socket;
    base::TimeTicks start_time;
    bool used;  // Indicates whether or not the socket has been used yet.

    // An idle socket should be removed if it can't be reused, or has been idle
    // for too long. |now| is the current time value (TimeTicks::Now()).
    // |timeout| is the length of time to wait before timing out an idle socket.
    //
    // An idle socket can't be reused if it is disconnected or has received
    // data unexpectedly (hence no longer idle).  The unread data would be
    // mistaken for the beginning of the next response if we were to reuse the
    // socket for a new request.
    bool ShouldCleanup(base::TimeTicks now, base::TimeDelta timeout) const;
  };

  typedef std::deque<const Request*> RequestQueue;
  typedef std::map<const ClientSocketHandle*, const Request*> RequestMap;

  // A Group is allocated per group_name when there are idle sockets or pending
  // requests.  Otherwise, the Group object is removed from the map.
  // |active_socket_count| tracks the number of sockets held by clients.
  struct Group {
    Group()
        : active_socket_count(0),
          backup_job(NULL),
          backup_task(NULL) {
    }

    ~Group() {
      CleanupBackupJob();
    }

    bool IsEmpty() const {
      return active_socket_count == 0 && idle_sockets.empty() && jobs.empty() &&
          pending_requests.empty();
    }

    bool HasAvailableSocketSlot(int max_sockets_per_group) const {
      return active_socket_count + static_cast<int>(jobs.size()) <
          max_sockets_per_group;
    }

    bool IsStalled(int max_sockets_per_group) const {
      return HasAvailableSocketSlot(max_sockets_per_group) &&
          pending_requests.size() > jobs.size();
    }

    RequestPriority TopPendingPriority() const {
      return pending_requests.front()->priority();
    }

    void CleanupBackupJob() {
      if (backup_job) {
        delete backup_job;
        backup_job = NULL;
      }
      if (backup_task) {
        backup_task->Cancel();
        backup_task = NULL;
      }
    }

    std::deque<IdleSocket> idle_sockets;
    std::set<const ConnectJob*> jobs;
    RequestQueue pending_requests;
    int active_socket_count;  // number of active sockets used by clients
    // A backup job in case the connect for this group takes too long.
    ConnectJob* backup_job;
    CancelableTask* backup_task;
  };

  typedef std::map<std::string, Group> GroupMap;

  typedef std::set<const ConnectJob*> ConnectJobSet;

  struct CallbackResultPair {
    CallbackResultPair() : callback(NULL), result(OK) {}
    CallbackResultPair(CompletionCallback* callback_in, int result_in)
        : callback(callback_in), result(result_in) {}

    CompletionCallback* callback;
    int result;
  };

  typedef std::map<const ClientSocketHandle*, CallbackResultPair>
      PendingCallbackMap;

  ~ClientSocketPoolBaseHelper();

  static void InsertRequestIntoQueue(const Request* r,
                                     RequestQueue* pending_requests);
  static const Request* RemoveRequestFromQueue(RequestQueue::iterator it,
                                               RequestQueue* pending_requests);

  // Called when the number of idle sockets changes.
  void IncrementIdleCount();
  void DecrementIdleCount();

  // Scans the group map for groups which have an available socket slot and
  // at least one pending request. Returns true if any groups are stalled, and
  // if so, fills |group| and |group_name| with data of the stalled group
  // having highest priority.
  bool FindTopStalledGroup(Group** group, std::string* group_name);

  // Called when timer_ fires.  This method scans the idle sockets removing
  // sockets that timed out or can't be reused.
  void OnCleanupTimerFired() {
    CleanupIdleSockets(false);
  }

  // Removes |job| from |connect_job_set_|.  Also updates |group| if non-NULL.
  void RemoveConnectJob(const ConnectJob* job, Group* group);

  // Tries to see if we can handle any more requests for |group|.
  void OnAvailableSocketSlot(const std::string& group_name, Group* group);

  // Process a pending socket request for a group.
  void ProcessPendingRequest(const std::string& group_name, Group* group);

  // Assigns |socket| to |handle| and updates |group|'s counters appropriately.
  void HandOutSocket(ClientSocket* socket,
                     bool reused,
                     ClientSocketHandle* handle,
                     base::TimeDelta time_idle,
                     Group* group,
                     const BoundNetLog& net_log);

  // Adds |socket| to the list of idle sockets for |group|.  |used| indicates
  // whether or not the socket has previously been used.
  void AddIdleSocket(ClientSocket* socket, bool used, Group* group);

  // Iterates through |connect_job_map_|, canceling all ConnectJobs.
  // Afterwards, it iterates through all groups and deletes them if they are no
  // longer needed.
  void CancelAllConnectJobs();

  // Returns true if we can't create any more sockets due to the total limit.
  bool ReachedMaxSocketsLimit() const;

  // This is the internal implementation of RequestSocket().  It differs in that
  // it does not handle logging into NetLog of the queueing status of
  // |request|.
  int RequestSocketInternal(const std::string& group_name,
                            const Request* request);

  // Assigns an idle socket for the group to the request.
  // Returns |true| if an idle socket is available, false otherwise.
  bool AssignIdleSocketToGroup(Group* group, const Request* request);

  static void LogBoundConnectJobToRequest(
      const NetLog::Source& connect_job_source, const Request* request);

  // Set a timer to create a backup socket if it takes too long to create one.
  void StartBackupSocketTimer(const std::string& group_name);

  // Called when the backup socket timer fires.
  void OnBackupSocketTimerFired(const std::string& group_name);

  // Closes one idle socket.  Picks the first one encountered.
  // TODO(willchan): Consider a better algorithm for doing this.  Perhaps we
  // should keep an ordered list of idle sockets, and close them in order.
  // Requires maintaining more state.  It's not clear if it's worth it since
  // I'm not sure if we hit this situation often.
  void CloseOneIdleSocket();

  // Checks if there are stalled socket groups that should be notified
  // for possible wakeup.
  void CheckForStalledSocketGroups();

  // Posts a task to call InvokeUserCallback() on the next iteration through the
  // current message loop.  Inserts |callback| into |pending_callback_map_|,
  // keyed by |handle|.
  void InvokeUserCallbackLater(
      ClientSocketHandle* handle, CompletionCallback* callback, int rv);

  // Invokes the user callback for |handle|.  By the time this task has run,
  // it's possible that the request has been cancelled, so |handle| may not
  // exist in |pending_callback_map_|.  We look up the callback and result code
  // in |pending_callback_map_|.
  void InvokeUserCallback(ClientSocketHandle* handle);

  GroupMap group_map_;

  // Map of the ClientSocketHandles for which we have a pending Task to invoke a
  // callback.  This is necessary since, before we invoke said callback, it's
  // possible that the request is cancelled.
  PendingCallbackMap pending_callback_map_;

  // Timer used to periodically prune idle sockets that timed out or can't be
  // reused.
  base::RepeatingTimer<ClientSocketPoolBaseHelper> timer_;

  // The total number of idle sockets in the system.
  int idle_socket_count_;

  // Number of connecting sockets across all groups.
  int connecting_socket_count_;

  // Number of connected sockets we handed out across all groups.
  int handed_out_socket_count_;

  // The maximum total number of sockets. See ReachedMaxSocketsLimit.
  const int max_sockets_;

  // The maximum number of sockets kept per group.
  const int max_sockets_per_group_;

  // The time to wait until closing idle sockets.
  const base::TimeDelta unused_idle_socket_timeout_;
  const base::TimeDelta used_idle_socket_timeout_;

  const scoped_ptr<ConnectJobFactory> connect_job_factory_;

  // TODO(vandebo) Remove when backup jobs move to TCPClientSocketPool
  bool backup_jobs_enabled_;

  // A factory to pin the backup_job tasks.
  ScopedRunnableMethodFactory<ClientSocketPoolBaseHelper> method_factory_;

  // A unique id for the pool.  It gets incremented every time we Flush() the
  // pool.  This is so that when sockets get released back to the pool, we can
  // make sure that they are discarded rather than reused.
  int pool_generation_number_;

  // Some parts of this class need to know if the destructor is running.
  bool in_destructor_;
};

}  // namespace internal

// The maximum duration, in seconds, to keep used idle persistent sockets alive.
static const int kUsedIdleSocketTimeout = 300;  // 5 minutes

template <typename SocketParams>
class ClientSocketPoolBase {
 public:
  class Request : public internal::ClientSocketPoolBaseHelper::Request {
   public:
    Request(ClientSocketHandle* handle,
            CompletionCallback* callback,
            RequestPriority priority,
            const scoped_refptr<SocketParams>& params,
            const BoundNetLog& net_log)
        : internal::ClientSocketPoolBaseHelper::Request(
              handle, callback, priority, net_log),
          params_(params) {}

    const scoped_refptr<SocketParams>& params() const { return params_; }

   private:
    scoped_refptr<SocketParams> params_;
  };

  class ConnectJobFactory {
   public:
    ConnectJobFactory() {}
    virtual ~ConnectJobFactory() {}

    virtual ConnectJob* NewConnectJob(
        const std::string& group_name,
        const Request& request,
        ConnectJob::Delegate* delegate) const = 0;

    virtual base::TimeDelta ConnectionTimeout() const = 0;

   private:
    DISALLOW_COPY_AND_ASSIGN(ConnectJobFactory);
  };

  // |max_sockets| is the maximum number of sockets to be maintained by this
  // ClientSocketPool.  |max_sockets_per_group| specifies the maximum number of
  // sockets a "group" can have.  |unused_idle_socket_timeout| specifies how
  // long to leave an unused idle socket open before closing it.
  // |used_idle_socket_timeout| specifies how long to leave a previously used
  // idle socket open before closing it.
  ClientSocketPoolBase(
      int max_sockets,
      int max_sockets_per_group,
      const scoped_refptr<ClientSocketPoolHistograms>& histograms,
      base::TimeDelta unused_idle_socket_timeout,
      base::TimeDelta used_idle_socket_timeout,
      ConnectJobFactory* connect_job_factory)
      : histograms_(histograms),
        helper_(new internal::ClientSocketPoolBaseHelper(
            max_sockets, max_sockets_per_group,
            unused_idle_socket_timeout, used_idle_socket_timeout,
            new ConnectJobFactoryAdaptor(connect_job_factory))) {}

  virtual ~ClientSocketPoolBase() {}

  // These member functions simply forward to ClientSocketPoolBaseHelper.

  // RequestSocket bundles up the parameters into a Request and then forwards to
  // ClientSocketPoolBaseHelper::RequestSocket().  Note that the memory
  // ownership is transferred in the asynchronous (ERR_IO_PENDING) case.
  int RequestSocket(const std::string& group_name,
                    const scoped_refptr<SocketParams>& params,
                    RequestPriority priority,
                    ClientSocketHandle* handle,
                    CompletionCallback* callback,
                    const BoundNetLog& net_log) {
    Request* request = new Request(handle, callback, priority, params, net_log);
    return helper_->RequestSocket(group_name, request);
  }

  void CancelRequest(const std::string& group_name,
                     ClientSocketHandle* handle) {
    return helper_->CancelRequest(group_name, handle);
  }

  void ReleaseSocket(const std::string& group_name, ClientSocket* socket,
                     int id) {
    return helper_->ReleaseSocket(group_name, socket, id);
  }

  void CloseIdleSockets() { return helper_->CloseIdleSockets(); }

  int idle_socket_count() const { return helper_->idle_socket_count(); }

  int IdleSocketCountInGroup(const std::string& group_name) const {
    return helper_->IdleSocketCountInGroup(group_name);
  }

  LoadState GetLoadState(const std::string& group_name,
                         const ClientSocketHandle* handle) const {
    return helper_->GetLoadState(group_name, handle);
  }

  virtual void OnConnectJobComplete(int result, ConnectJob* job) {
    return helper_->OnConnectJobComplete(result, job);
  }

  int NumConnectJobsInGroup(const std::string& group_name) const {
    return helper_->NumConnectJobsInGroup(group_name);
  }

  void CleanupIdleSockets(bool force) {
    return helper_->CleanupIdleSockets(force);
  }

  base::TimeDelta ConnectionTimeout() const {
    return helper_->ConnectionTimeout();
  }

  scoped_refptr<ClientSocketPoolHistograms> histograms() const {
    return histograms_;
  }

  void EnableBackupJobs() { helper_->EnableBackupJobs(); }

  void Flush() { helper_->Flush(); }

 private:
  // This adaptor class exists to bridge the
  // internal::ClientSocketPoolBaseHelper::ConnectJobFactory and
  // ClientSocketPoolBase::ConnectJobFactory types, allowing clients to use the
  // typesafe ClientSocketPoolBase::ConnectJobFactory, rather than having to
  // static_cast themselves.
  class ConnectJobFactoryAdaptor
      : public internal::ClientSocketPoolBaseHelper::ConnectJobFactory {
   public:
    typedef typename ClientSocketPoolBase<SocketParams>::ConnectJobFactory
        ConnectJobFactory;

    explicit ConnectJobFactoryAdaptor(ConnectJobFactory* connect_job_factory)
        : connect_job_factory_(connect_job_factory) {}
    virtual ~ConnectJobFactoryAdaptor() {}

    virtual ConnectJob* NewConnectJob(
        const std::string& group_name,
        const internal::ClientSocketPoolBaseHelper::Request& request,
        ConnectJob::Delegate* delegate) const {
      const Request* casted_request = static_cast<const Request*>(&request);
      return connect_job_factory_->NewConnectJob(
          group_name, *casted_request, delegate);
    }

    virtual base::TimeDelta ConnectionTimeout() const {
      return connect_job_factory_->ConnectionTimeout();
    }

    const scoped_ptr<ConnectJobFactory> connect_job_factory_;
  };

  // Histograms for the pool
  const scoped_refptr<ClientSocketPoolHistograms> histograms_;

  // The reason for reference counting here is because the operations on
  // the ClientSocketPoolBaseHelper which release sockets can cause the
  // ClientSocketPoolBase<T> reference to drop to zero.  While we're deep
  // in cleanup code, we'll often hold a reference to |self|.
  scoped_refptr<internal::ClientSocketPoolBaseHelper> helper_;

  DISALLOW_COPY_AND_ASSIGN(ClientSocketPoolBase);
};

}  // namespace net

#endif  // NET_SOCKET_CLIENT_SOCKET_POOL_BASE_H_