1 // Copyright 2023 The Chromium Authors. All rights reserved.
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4 
5 #ifndef QUICHE_WEB_TRANSPORT_WEB_TRANSPORT_PRIORITY_SCHEDULER_H_
6 #define QUICHE_WEB_TRANSPORT_WEB_TRANSPORT_PRIORITY_SCHEDULER_H_
7 
8 #include <cstddef>
9 #include <optional>
10 #include <utility>
11 
12 #include "absl/container/flat_hash_map.h"
13 #include "absl/container/node_hash_map.h"
14 #include "absl/status/status.h"
15 #include "absl/status/statusor.h"
16 #include "quiche/common/btree_scheduler.h"
17 #include "quiche/common/platform/api/quiche_export.h"
18 #include "quiche/web_transport/web_transport.h"
19 
20 namespace webtransport {
21 
22 // Schedules the streams within a WebTransport session as defined by the W3C
23 // WebTransport API. Unlike the W3C API, there is no need to track groups
24 // manually: the group is created when a first stream with the associated group
25 // ID is registered, and it is deleted when the last stream associated with the
26 // group is unregistered.
27 class QUICHE_EXPORT PriorityScheduler {
28  public:
29   // Returns true if there are any streams registered.
HasRegistered()30   bool HasRegistered() const { return active_groups_.HasRegistered(); }
31   // Returns true if there are any streams scheduled.
HasScheduled()32   bool HasScheduled() const { return active_groups_.HasScheduled(); }
33 
34   // Returns the number of currently scheduled streams.
NumScheduled()35   size_t NumScheduled() const {
36     size_t total = 0;
37     for (const auto& [group_id, group_scheduler] : per_group_schedulers_) {
38       total += group_scheduler.NumScheduled();
39     }
40     return total;
41   }
42 
43   // Registers the specified stream with the supplied priority.  The stream must
44   // not be already registered.
45   absl::Status Register(StreamId stream_id, const StreamPriority& priority);
46   // Unregisters a previously registered stream.
47   absl::Status Unregister(StreamId stream_id);
48   // Alters the priority of an already registered stream.
49   absl::Status UpdateSendOrder(StreamId stream_id, SendOrder new_send_order);
50   absl::Status UpdateSendGroup(StreamId stream_id, SendGroupId new_send_group);
51 
52   // Returns true if there is a stream that would go before `id` in the
53   // schedule.
54   absl::StatusOr<bool> ShouldYield(StreamId id) const;
55 
56   // Returns the priority for `id`, or nullopt if stream is not registered.
57   std::optional<StreamPriority> GetPriorityFor(StreamId stream_id) const;
58 
59   // Pops the highest priority stream.  Will fail if the schedule is empty.
60   absl::StatusOr<StreamId> PopFront();
61 
62   // Adds `stream` to the schedule if it's not already there.
63   absl::Status Schedule(StreamId stream_id);
64   // Returns true if `stream` is in the schedule.
65   bool IsScheduled(StreamId stream_id) const;
66 
67  private:
68   // All groups currently have the equal priority; this type represents the said
69   // single priority.
70   class SinglePriority {
71    public:
72     bool operator==(const SinglePriority&) const { return true; }
73     bool operator!=(const SinglePriority&) const { return false; }
74 
75     bool operator<(const SinglePriority&) const { return false; }
76     bool operator>(const SinglePriority&) const { return false; }
77     bool operator<=(const SinglePriority&) const { return true; }
78     bool operator>=(const SinglePriority&) const { return true; }
79   };
80 
81   using PerGroupScheduler = quiche::BTreeScheduler<StreamId, SendOrder>;
82   using GroupSchedulerPair = std::pair<const SendGroupId, PerGroupScheduler>;
83 
84   // Round-robin schedule for the groups.
85   quiche::BTreeScheduler<SendGroupId, SinglePriority> active_groups_;
86   // Map group ID to the scheduler for the group in question.
87   absl::node_hash_map<SendGroupId, PerGroupScheduler> per_group_schedulers_;
88   // Map stream ID to a pointer to the entry in `per_group_schedulers_`.
89   absl::flat_hash_map<StreamId, GroupSchedulerPair*> stream_to_group_map_;
90 
SchedulerForStream(StreamId stream_id)91   PerGroupScheduler* SchedulerForStream(StreamId stream_id) {
92     auto it = stream_to_group_map_.find(stream_id);
93     if (it == stream_to_group_map_.end()) {
94       return nullptr;
95     }
96     return &it->second->second;
97   }
SchedulerForStream(StreamId stream_id)98   const PerGroupScheduler* SchedulerForStream(StreamId stream_id) const {
99     auto it = stream_to_group_map_.find(stream_id);
100     if (it == stream_to_group_map_.end()) {
101       return nullptr;
102     }
103     return &it->second->second;
104   }
105 };
106 
107 }  // namespace webtransport
108 
109 #endif  // QUICHE_WEB_TRANSPORT_WEB_TRANSPORT_PRIORITY_SCHEDULER_H_
110