1// Copyright (c) 2013 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 NET_SPDY_SPDY_PRIORITY_FOREST_H_
6#define NET_SPDY_SPDY_PRIORITY_FOREST_H_
7
8#include <map>
9#include <set>
10#include <vector>
11
12#include "base/basictypes.h"
13#include "base/containers/hash_tables.h"
14#include "base/logging.h"
15#include "base/memory/scoped_ptr.h"
16#include "base/rand_util.h"
17
18namespace net {
19
20// This data structure implements the SPDY prioriziation data structures
21// defined in this document: http://go/spdy4-prioritization
22//
23// Nodes can be added and removed, and dependencies between them defined.  Each
24// node can have at most one parent and at most one child (forming a list), but
25// there can be multiple lists, with each list root having its own priority.
26// Individual nodes can also be marked as ready to read/write, and then the
27// whole structure can be queried to pick the next node to read/write out of
28// those ready.
29//
30// The NodeId and Priority types must be POD that support comparison (most
31// likely, they will be numbers).
32template <typename NodeId, typename Priority>
33class SpdyPriorityForest {
34 public:
35  SpdyPriorityForest();
36  ~SpdyPriorityForest();
37
38  // Return the number of nodes currently in the forest.
39  int num_nodes() const;
40
41  // Return true if the forest contains a node with the given ID.
42  bool NodeExists(NodeId node_id) const;
43
44  // Add a new root node to the forest, with the given priority.  Returns true
45  // on success, or false if the node_id already exists within the forest.
46  bool AddRootNode(NodeId node_id, Priority priority);
47
48  // Add a new node to the forest, with the given parent.  Returns true on
49  // success.  Returns false and has no effect if the new node already exists,
50  // or if the parent doesn't exist, or if the parent already has a child.
51  bool AddNonRootNode(NodeId node_id, NodeId parent_id, bool unordered);
52
53  // Remove an existing node from the forest.  Returns true on success, or
54  // false if the node doesn't exist.
55  bool RemoveNode(NodeId node_id);
56
57  // Get the priority of the given node.  If the node doesn't exist, or is not
58  // a root node (and thus has no priority), returns Priority().
59  Priority GetPriority(NodeId node_id) const;
60
61  // Get the parent of the given node.  If the node doesn't exist, or is a root
62  // node (and thus has no parent), returns NodeId().
63  NodeId GetParent(NodeId node_id) const;
64
65  // Determine if the given node is unordered with respect to its parent.  If
66  // the node doesn't exist, or is a root node (and thus has no parent),
67  // returns false.
68  bool IsNodeUnordered(NodeId node_id) const;
69
70  // Get the child of the given node.  If the node doesn't exist, or has no
71  // child, returns NodeId().
72  NodeId GetChild(NodeId node_id) const;
73
74  // Set the priority of the given node.  If the node was not already a root
75  // node, this makes it a root node.  Returns true on success, or false if the
76  // node doesn't exist.
77  bool SetPriority(NodeId node_id, Priority priority);
78
79  // Set the parent of the given node.  If the node was a root node, this makes
80  // it no longer a root.  Returns true on success.  Returns false and has no
81  // effect if (1) the node and/or the parent doesn't exist, (2) the new parent
82  // already has a different child than the node, or (3) if the new parent is a
83  // descendant of the node (so this would have created a cycle).
84  bool SetParent(NodeId node_id, NodeId parent_id, bool unordered);
85
86  // Check if a node is marked as ready to read.  Returns false if the node
87  // doesn't exist.
88  bool IsMarkedReadyToRead(NodeId node_id) const;
89  // Mark the node as ready or not ready to read.  Returns true on success, or
90  // false if the node doesn't exist.
91  bool MarkReadyToRead(NodeId node_id);
92  bool MarkNoLongerReadyToRead(NodeId node_id);
93  // Return the ID of the next node that we should read, or return NodeId() if
94  // no node in the forest is ready to read.
95  NodeId NextNodeToRead();
96
97  // Check if a node is marked as ready to write.  Returns false if the node
98  // doesn't exist.
99  bool IsMarkedReadyToWrite(NodeId node_id) const;
100  // Mark the node as ready or not ready to write.  Returns true on success, or
101  // false if the node doesn't exist.
102  bool MarkReadyToWrite(NodeId node_id);
103  bool MarkNoLongerReadyToWrite(NodeId node_id);
104  // Return the ID of the next node that we should write, or return NodeId() if
105  // no node in the forest is ready to write.
106  NodeId NextNodeToWrite();
107
108  // Return true if all internal invariants hold (useful for unit tests).
109  // Unless there are bugs, this should always return true.
110  bool ValidateInvariantsForTests() const;
111
112 private:
113  enum NodeType { ROOT_NODE, NONROOT_ORDERED, NONROOT_UNORDERED };
114  struct Node {
115    Node() : type(ROOT_NODE), flags(0), child() {
116      depends_on.priority = Priority();
117    }
118    NodeType type;
119    unsigned int flags;  // bitfield of flags
120    union {
121      Priority priority;  // used for root nodes
122      NodeId parent_id;  // used for non-root nodes
123    } depends_on;
124    NodeId child;  // node ID of child (or NodeId() for no child)
125  };
126
127  typedef base::hash_map<NodeId, Node> NodeMap;
128
129  // Constants for the Node.flags bitset:
130  // kReadToRead: set for nodes that are ready for reading
131  static const unsigned int kReadyToRead = (1 << 0);
132  // kReadToWrite: set for nodes that are ready for writing
133  static const unsigned int kReadyToWrite = (1 << 1);
134
135  // Common code for IsMarkedReadyToRead and IsMarkedReadyToWrite.
136  bool IsMarked(NodeId node_id, unsigned int flag) const;
137  // Common code for MarkReadyToRead and MarkReadyToWrite.
138  bool Mark(NodeId node_id, unsigned int flag);
139  // Common code for MarkNoLongerReadyToRead and MarkNoLongerReadyToWrite.
140  bool Unmark(NodeId node_id, unsigned int flag);
141  // Common code for NextNodeToRead and NextNodeToWrite;
142  NodeId FirstMarkedNode(unsigned int flag);
143  // Get the given node, or return NULL if it doesn't exist.
144  const Node* FindNode(NodeId node_id) const;
145
146  NodeMap all_nodes_;  // maps from node IDs to Node objects
147
148  DISALLOW_COPY_AND_ASSIGN(SpdyPriorityForest);
149};
150
151template <typename NodeId, typename Priority>
152SpdyPriorityForest<NodeId, Priority>::SpdyPriorityForest() {}
153
154template <typename NodeId, typename Priority>
155SpdyPriorityForest<NodeId, Priority>::~SpdyPriorityForest() {}
156
157template <typename NodeId, typename Priority>
158int SpdyPriorityForest<NodeId, Priority>::num_nodes() const {
159  return all_nodes_.size();
160}
161
162template <typename NodeId, typename Priority>
163bool SpdyPriorityForest<NodeId, Priority>::NodeExists(NodeId node_id) const {
164  return all_nodes_.count(node_id) != 0;
165}
166
167template <typename NodeId, typename Priority>
168bool SpdyPriorityForest<NodeId, Priority>::AddRootNode(
169    NodeId node_id, Priority priority) {
170  if (NodeExists(node_id)) {
171    return false;
172  }
173  Node* new_node = &all_nodes_[node_id];
174  new_node->type = ROOT_NODE;
175  new_node->depends_on.priority = priority;
176  return true;
177}
178
179template <typename NodeId, typename Priority>
180bool SpdyPriorityForest<NodeId, Priority>::AddNonRootNode(
181    NodeId node_id, NodeId parent_id, bool unordered) {
182  if (NodeExists(node_id) || !NodeExists(parent_id)) {
183    return false;
184  }
185
186  Node* parent = &all_nodes_[parent_id];
187  if (parent->child != NodeId()) {
188    return false;
189  }
190
191  Node* new_node = &all_nodes_[node_id];
192  new_node->type = (unordered ? NONROOT_UNORDERED : NONROOT_ORDERED);
193  new_node->depends_on.parent_id = parent_id;
194  parent->child = node_id;
195  return true;
196}
197
198template <typename NodeId, typename Priority>
199bool SpdyPriorityForest<NodeId, Priority>::RemoveNode(NodeId node_id) {
200  if (!NodeExists(node_id)) {
201    return false;
202  }
203  const Node& node = all_nodes_[node_id];
204
205  // If the node to be removed is not a root node, we need to change its
206  // parent's child ID.
207  if (node.type != ROOT_NODE) {
208    DCHECK(NodeExists(node.depends_on.parent_id));
209    Node* parent = &all_nodes_[node.depends_on.parent_id];
210    DCHECK_EQ(node_id, parent->child);
211    parent->child = node.child;
212  }
213
214  // If the node has a child, we need to change the child's priority or parent.
215  if (node.child != NodeId()) {
216    DCHECK(NodeExists(node.child));
217    Node* child = &all_nodes_[node.child];
218    DCHECK_NE(ROOT_NODE, child->type);
219    DCHECK_EQ(node_id, child->depends_on.parent_id);
220    // Make the child's new depends_on be the node's depends_on (whether that
221    // be a priority or a parent node ID).
222    child->depends_on = node.depends_on;
223    // If the removed node was a root, its child is now a root.  Otherwise, the
224    // child will be be unordered if and only if it was already unordered and
225    // the removed not is also not ordered.
226    if (node.type == ROOT_NODE) {
227      child->type = ROOT_NODE;
228    } else if (node.type == NONROOT_ORDERED) {
229      child->type = NONROOT_ORDERED;
230    }
231  }
232
233  // Delete the node.
234  all_nodes_.erase(node_id);
235  return true;
236}
237
238template <typename NodeId, typename Priority>
239Priority SpdyPriorityForest<NodeId, Priority>::GetPriority(
240    NodeId node_id) const {
241  const Node* node = FindNode(node_id);
242  if (node != NULL && node->type == ROOT_NODE) {
243    return node->depends_on.priority;
244  } else {
245    return Priority();
246  }
247}
248
249template <typename NodeId, typename Priority>
250NodeId SpdyPriorityForest<NodeId, Priority>::GetParent(NodeId node_id) const {
251  const Node* node = FindNode(node_id);
252  if (node != NULL && node->type != ROOT_NODE) {
253    return node->depends_on.parent_id;
254  } else {
255    return NodeId();
256  }
257}
258
259template <typename NodeId, typename Priority>
260bool SpdyPriorityForest<NodeId, Priority>::IsNodeUnordered(
261    NodeId node_id) const {
262  const Node* node = FindNode(node_id);
263  return node != NULL && node->type == NONROOT_UNORDERED;
264}
265
266template <typename NodeId, typename Priority>
267NodeId SpdyPriorityForest<NodeId, Priority>::GetChild(NodeId node_id) const {
268  const Node* node = FindNode(node_id);
269  if (node != NULL) {
270    return node->child;
271  } else {
272    return NodeId();
273  }
274}
275
276template <typename NodeId, typename Priority>
277bool SpdyPriorityForest<NodeId, Priority>::SetPriority(
278    NodeId node_id, Priority priority) {
279  if (!NodeExists(node_id)) {
280    return false;
281  }
282
283  Node* node = &all_nodes_[node_id];
284  // If this is not already a root node, we need to make it be a root node.
285  if (node->type != ROOT_NODE) {
286    DCHECK(NodeExists(node->depends_on.parent_id));
287    Node* parent = &all_nodes_[node->depends_on.parent_id];
288    parent->child = NodeId();
289    node->type = ROOT_NODE;
290  }
291
292  node->depends_on.priority = priority;
293  return true;
294}
295
296template <typename NodeId, typename Priority>
297bool SpdyPriorityForest<NodeId, Priority>::SetParent(
298    NodeId node_id, NodeId parent_id, bool unordered) {
299  if (!NodeExists(node_id) || !NodeExists(parent_id)) {
300    return false;
301  }
302
303  Node* node = &all_nodes_[node_id];
304  Node* new_parent = &all_nodes_[parent_id];
305  // If the new parent is already the node's parent, all we have to do is
306  // update the node type and we're done.
307  if (new_parent->child == node_id) {
308    node->type = (unordered ? NONROOT_UNORDERED : NONROOT_ORDERED);
309    return true;
310  }
311  // Otherwise, if the new parent already has a child, we fail.
312  if (new_parent->child != NodeId()) {
313    return false;
314  }
315
316  // Next, make sure we won't create a cycle.
317  if (node_id == parent_id) return false;
318  Node* last = node;
319  NodeId last_id = node_id;
320  while (last->child != NodeId()) {
321    if (last->child == parent_id) return false;
322    last_id = last->child;
323    DCHECK(NodeExists(last_id));
324    last = &all_nodes_[last_id];
325  }
326
327  // If the node is not a root, we need clear its old parent's child field
328  // (unless the old parent is the same as the new parent).
329  if (node->type != ROOT_NODE) {
330    const NodeId old_parent_id = node->depends_on.parent_id;
331    DCHECK(NodeExists(old_parent_id));
332    DCHECK(old_parent_id != parent_id);
333    Node* old_parent = &all_nodes_[old_parent_id];
334    DCHECK_EQ(node_id, old_parent->child);
335    old_parent->child = NodeId();
336  }
337
338  // Make the change.
339  node->type = (unordered ? NONROOT_UNORDERED : NONROOT_ORDERED);
340  node->depends_on.parent_id = parent_id;
341  new_parent->child = node_id;
342  return true;
343}
344
345template <typename NodeId, typename Priority>
346bool SpdyPriorityForest<NodeId, Priority>::IsMarkedReadyToRead(
347    NodeId node_id) const {
348  return IsMarked(node_id, kReadyToRead);
349}
350
351template <typename NodeId, typename Priority>
352bool SpdyPriorityForest<NodeId, Priority>::MarkReadyToRead(NodeId node_id) {
353  return Mark(node_id, kReadyToRead);
354}
355
356template <typename NodeId, typename Priority>
357bool SpdyPriorityForest<NodeId, Priority>::MarkNoLongerReadyToRead(
358    NodeId node_id) {
359  return Unmark(node_id, kReadyToRead);
360}
361
362template <typename NodeId, typename Priority>
363NodeId SpdyPriorityForest<NodeId, Priority>::NextNodeToRead() {
364  return FirstMarkedNode(kReadyToRead);
365}
366
367template <typename NodeId, typename Priority>
368bool SpdyPriorityForest<NodeId, Priority>::IsMarkedReadyToWrite(
369    NodeId node_id) const {
370  return IsMarked(node_id, kReadyToWrite);
371}
372
373template <typename NodeId, typename Priority>
374bool SpdyPriorityForest<NodeId, Priority>::MarkReadyToWrite(NodeId node_id) {
375  return Mark(node_id, kReadyToWrite);
376}
377
378template <typename NodeId, typename Priority>
379bool SpdyPriorityForest<NodeId, Priority>::MarkNoLongerReadyToWrite(
380    NodeId node_id) {
381  return Unmark(node_id, kReadyToWrite);
382}
383
384template <typename NodeId, typename Priority>
385NodeId SpdyPriorityForest<NodeId, Priority>::NextNodeToWrite() {
386  return FirstMarkedNode(kReadyToWrite);
387}
388
389template <typename NodeId, typename Priority>
390bool SpdyPriorityForest<NodeId, Priority>::IsMarked(
391    NodeId node_id, unsigned int flag) const {
392  const Node* node = FindNode(node_id);
393  return node != NULL && (node->flags & flag) != 0;
394}
395
396template <typename NodeId, typename Priority>
397bool SpdyPriorityForest<NodeId, Priority>::Mark(
398    NodeId node_id, unsigned int flag) {
399  if (!NodeExists(node_id)) {
400    return false;
401  }
402  all_nodes_[node_id].flags |= flag;
403  return true;
404}
405
406template <typename NodeId, typename Priority>
407bool SpdyPriorityForest<NodeId, Priority>::Unmark(
408    NodeId node_id, unsigned int flag) {
409  if (!NodeExists(node_id)) {
410    return false;
411  }
412  all_nodes_[node_id].flags &= ~flag;
413  return true;
414}
415
416template <typename NodeId, typename Priority>
417NodeId SpdyPriorityForest<NodeId, Priority>::FirstMarkedNode(
418    unsigned int flag) {
419  // TODO(mdsteele): This is an *incredibly* stupid brute force solution.
420
421  // Get all root nodes that have at least one marked child.
422  uint64 total_weight = 0;
423  std::map<uint64, NodeId> roots;  // maps cumulative weight to root node ID
424  for (typename NodeMap::const_iterator iter = all_nodes_.begin();
425       iter != all_nodes_.end(); ++iter) {
426    const NodeId root_id = iter->first;
427    const Node& root = iter->second;
428    if (root.type == ROOT_NODE) {
429      // See if there is at least one marked node in this root's chain.
430      for (const Node* node = &root; ; node = &all_nodes_[node->child]) {
431        if ((node->flags & flag) != 0) {
432          total_weight += static_cast<uint64>(root.depends_on.priority);
433          roots[total_weight] = root_id;
434          break;
435        }
436        if (node->child == NodeId()) {
437          break;
438        }
439        DCHECK(NodeExists(node->child));
440      }
441    }
442  }
443
444  // If there are no ready nodes, then return NodeId().
445  if (total_weight == 0) {
446    DCHECK(roots.empty());
447    return NodeId();
448  } else {
449    DCHECK(!roots.empty());
450  }
451
452  // Randomly select a tree to use.
453  typename std::map<uint64, NodeId>::const_iterator root_iter =
454      roots.upper_bound(base::RandGenerator(total_weight));
455  DCHECK(root_iter != roots.end());
456  const NodeId root_id = root_iter->second;
457
458  // Find the first node in the chain that is ready.
459  NodeId node_id = root_id;
460  while (true) {
461    DCHECK(NodeExists(node_id));
462    Node* node = &all_nodes_[node_id];
463    if ((node->flags & flag) != 0) {
464      // There might be more nodes that are ready and that are linked to this
465      // one in an unordered chain.  Find all of them, then pick one randomly.
466      std::vector<NodeId> group;
467      group.push_back(node_id);
468      for (Node* next = node; next->child != NodeId();) {
469        DCHECK(NodeExists(next->child));
470        Node *child = &all_nodes_[next->child];
471        DCHECK_NE(ROOT_NODE, child->type);
472        if (child->type != NONROOT_UNORDERED) {
473          break;
474        }
475        if ((child->flags & flag) != 0) {
476          group.push_back(next->child);
477        }
478        next = child;
479      }
480      return group[base::RandGenerator(group.size())];
481    }
482    node_id = node->child;
483  }
484}
485
486template <typename NodeId, typename Priority>
487const typename SpdyPriorityForest<NodeId, Priority>::Node*
488SpdyPriorityForest<NodeId, Priority>::FindNode(NodeId node_id) const {
489  typename NodeMap::const_iterator iter = all_nodes_.find(node_id);
490  if (iter == all_nodes_.end()) {
491    return NULL;
492  }
493  return &iter->second;
494}
495
496template <typename NodeId, typename Priority>
497bool SpdyPriorityForest<NodeId, Priority>::ValidateInvariantsForTests() const {
498  for (typename NodeMap::const_iterator iter = all_nodes_.begin();
499       iter != all_nodes_.end(); ++iter) {
500    const NodeId node_id = iter->first;
501    const Node& node = iter->second;
502    if (node.type != ROOT_NODE &&
503        (!NodeExists(node.depends_on.parent_id) ||
504         GetChild(node.depends_on.parent_id) != node_id)) {
505      return false;
506    }
507    if (node.child != NodeId()) {
508      if (!NodeExists(node.child) || node_id != GetParent(node.child)) {
509        return false;
510      }
511    }
512
513    NodeId child_id = node.child;
514    int count = 0;
515    while (child_id != NodeId()) {
516      if (count > num_nodes() || node_id == child_id) {
517        return false;
518      }
519      child_id = GetChild(child_id);
520      ++count;
521    }
522  }
523  return true;
524}
525
526}  // namespace net
527
528#endif  // NET_SPDY_SPDY_PRIORITY_FOREST_H_
529