// Copyright 2022 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.

#ifndef QUICHE_QUIC_CORE_IO_QUIC_POLL_EVENT_LOOP_H_
#define QUICHE_QUIC_CORE_IO_QUIC_POLL_EVENT_LOOP_H_

#include <poll.h>

#include <memory>

#include "absl/container/btree_map.h"
#include "absl/types/span.h"
#include "quiche/quic/core/io/quic_event_loop.h"
#include "quiche/quic/core/quic_alarm.h"
#include "quiche/quic/core/quic_alarm_factory.h"
#include "quiche/quic/core/quic_clock.h"
#include "quiche/quic/core/quic_udp_socket.h"
#include "quiche/common/quiche_linked_hash_map.h"

namespace quic {

// A simple and portable implementation of QuicEventLoop using poll(2).  Works
// on all POSIX platforms (and can be potentially made to support Windows using
// WSAPoll).
//
// For most operations, this implementation has a typical runtime of
// O(N + log M), where N is the number of file descriptors, and M is the number
// of pending alarms.
//
// This API has to deal with the situations where callbacks are modified from
// the callbacks themselves.  To address this, we use the following two
// approaches:
//   1. The code does not execute any callbacks until the very end of the
//      processing, when all of the state for the event loop is consistent.
//   2. The callbacks are stored as weak pointers, since other callbacks can
//      cause them to be unregistered.
class QUICHE_NO_EXPORT QuicPollEventLoop : public QuicEventLoop {
 public:
  QuicPollEventLoop(QuicClock* clock);

  // QuicEventLoop implementation.
  bool SupportsEdgeTriggered() const override { return false; }
  ABSL_MUST_USE_RESULT bool RegisterSocket(
      QuicUdpSocketFd fd, QuicSocketEventMask events,
      QuicSocketEventListener* listener) override;
  ABSL_MUST_USE_RESULT bool UnregisterSocket(QuicUdpSocketFd fd) override;
  ABSL_MUST_USE_RESULT bool RearmSocket(QuicUdpSocketFd fd,
                                        QuicSocketEventMask events) override;
  ABSL_MUST_USE_RESULT bool ArtificiallyNotifyEvent(
      QuicUdpSocketFd fd, QuicSocketEventMask events) override;
  void RunEventLoopOnce(QuicTime::Delta default_timeout) override;
  QuicAlarmFactory* GetAlarmFactory() override;

 protected:
  // Allows poll(2) calls to be mocked out in unit tests.
  virtual int PollSyscall(pollfd* fds, nfds_t nfds, int timeout) {
    return ::poll(fds, nfds, timeout);
  }

 private:
  friend class QuicPollEventLoopPeer;

  struct Registration {
    QuicSocketEventMask events = 0;
    QuicSocketEventListener* listener;

    QuicSocketEventMask artificially_notify_at_next_iteration = 0;
  };

  class Alarm : public QuicAlarm {
   public:
    Alarm(QuicPollEventLoop* loop,
          QuicArenaScopedPtr<QuicAlarm::Delegate> delegate);

    void SetImpl() override;
    void CancelImpl() override;

    void DoFire() {
      current_schedule_handle_.reset();
      Fire();
    }

   private:
    QuicPollEventLoop* loop_;
    // Deleted when the alarm is cancelled, causing the corresponding weak_ptr
    // in the alarm list to not be executed.
    std::shared_ptr<Alarm*> current_schedule_handle_;
  };

  class AlarmFactory : public QuicAlarmFactory {
   public:
    AlarmFactory(QuicPollEventLoop* loop) : loop_(loop) {}

    // QuicAlarmFactory implementation.
    QuicAlarm* CreateAlarm(QuicAlarm::Delegate* delegate) override;
    QuicArenaScopedPtr<QuicAlarm> CreateAlarm(
        QuicArenaScopedPtr<QuicAlarm::Delegate> delegate,
        QuicConnectionArena* arena) override;

   private:
    QuicPollEventLoop* loop_;
  };

  // Used for deferred execution of I/O callbacks.
  struct ReadyListEntry {
    QuicUdpSocketFd fd;
    std::weak_ptr<Registration> registration;
    QuicSocketEventMask events;
  };

  // We're using a linked hash map here to ensure the events are called in the
  // registration order.  This isn't strictly speaking necessary, but makes
  // testing things easier.
  using RegistrationMap =
      quiche::QuicheLinkedHashMap<QuicUdpSocketFd,
                                  std::shared_ptr<Registration>>;
  // Alarms are stored as weak pointers, since the alarm can be cancelled and
  // disappear while in the queue.
  using AlarmList = absl::btree_multimap<QuicTime, std::weak_ptr<Alarm*>>;

  // Returns the timeout for the next poll(2) call.  It is typically the time at
  // which the next alarm is supposed to activate.
  QuicTime::Delta ComputePollTimeout(QuicTime now,
                                     QuicTime::Delta default_timeout) const;
  // Calls poll(2) with the provided timeout and dispatches the callbacks
  // accordingly.
  void ProcessIoEvents(QuicTime start_time, QuicTime::Delta timeout);
  // Calls all of the alarm callbacks that are scheduled before or at |time|.
  void ProcessAlarmsUpTo(QuicTime time);

  // Adds the I/O callbacks for |fd| to the |ready_lits| as appopriate.
  void DispatchIoEvent(std::vector<ReadyListEntry>& ready_list,
                       QuicUdpSocketFd fd, short mask);  // NOLINT(runtime/int)
  // Runs all of the callbacks on the ready list.
  void RunReadyCallbacks(std::vector<ReadyListEntry>& ready_list);

  // Calls poll() while handling EINTR.  Returns the return value of poll(2)
  // system call.
  int PollWithRetries(absl::Span<pollfd> fds, QuicTime start_time,
                      QuicTime::Delta timeout);

  const QuicClock* clock_;
  RegistrationMap registrations_;
  AlarmList alarms_;
  AlarmFactory alarm_factory_;
  bool has_artificial_events_pending_ = false;
};

}  // namespace quic

#endif  // QUICHE_QUIC_CORE_IO_QUIC_POLL_EVENT_LOOP_H_