Improved the correctness of TCP (including both Reno and CUBIC variants) by refactoring the current congestion control codebase.

- Rebuilt the congestion-control base into LossBasedCongestionControl,
adding shared RFC5681 slow-start/fast-recovery handling,
cwnd-in-segments helpers, and loss/event hooks that Reno/CUBIC (and
future variants) can override without duplicating bookkeeping
(ns/flow/cc.py). TCPReno now simply supplies its congestion-avoidance
increment so the standard behavior stays centralized.
- Re-implemented CUBIC atop the new base with configurable beta/C,
proper epoch resets, fast-convergence bookkeeping, TCP-friendly
fallback, and persistent d_min tracking so its window evolution now
matches RFC8312 (§4) semantics (ns/flow/cubic.py).
- Added regression coverage for Reno slow start and fast recovery as
well as CUBIC’s loss/timeout response and cubic growth, plus a
tests/conftest.py shim so uv run pytest can import the in-tree ns
package without needing an install step
(tests/flow/test_tcp_congestion.py, tests/conftest.py).

Author: baochunli

ns/flow/cc.py

@@ -1,23 +1,36 @@
 """
-The base class for congestion control algorithms, designed to supply the TCPPacketGenerator class
-with congestion control decisions.
+Shared congestion-control infrastructure for loss-based TCP variants.
+
+The Simulator models cwnd in *bytes* so the helpers provided here enforce RFC 5681/8312
+requirements while letting individual algorithms focus on their window update rules.
 """
+from __future__ import annotations
+
 from abc import abstractmethod
+from enum import Enum, auto
+from typing import Final
+
+
+class LossEvent(Enum):
+    """Enumerates the two canonical loss signals in TCP."""
+
+    FAST_LOSS = auto()
+    TIMEOUT = auto()
 
 
 class CongestionControl:
     """
-    The base class for congestion control algorithms, designed to supply the TCPPacketGenerator
-    class with congestion control decisions.
+    Base class for congestion control algorithms, designed to supply TCPPacketGenerator
+    with congestion-control decisions.
 
     Parameters
     ----------
     mss: int
-        the maximum segment size
+        Maximum segment size in bytes.
     cwnd: int
-        the size of the congestion window.
+        Congestion window in bytes.
     ssthresh: int
-        the slow start threshold.
+        Slow-start threshold in bytes.
     debug: bool
         If True, prints more verbose debug information.
     """
@@ -47,38 +60,113 @@ def ack_received(self, rtt: float = 0, current_time: float = 0):
 
     def timer_expired(self, packet=None):
         """Actions to be taken when a timer expired."""
-        self.ssthresh = max(2 * self.mss, self.cwnd / 2)
-        # setting the congestion window to 1 segment
-        self.cwnd = self.mss
+        raise NotImplementedError("timer_expired must be implemented by subclasses.")
 
     def dupack_over(self):
         """Actions to be taken when a new ack is received after previous dupacks."""
-        # RFC 2001 and TCP Reno
-        self.cwnd = self.ssthresh
+        raise NotImplementedError("dupack_over must be implemented by subclasses.")
 
     def consecutive_dupacks_received(self, packet=None):
         """Actions to be taken when three consecutive dupacks are received."""
-        # fast retransmit in RFC 2001 and TCP Reno
-        self.ssthresh = max(2 * self.mss, self.cwnd / 2)
-        self.cwnd = self.ssthresh + 3 * self.mss
+        raise NotImplementedError(
+            "consecutive_dupacks_received must be implemented by subclasses."
+        )
 
     def more_dupacks_received(self, packet=None):
         """Actions to be taken when more than three consecutive dupacks are received."""
-        # fast retransmit in RFC 2001 and TCP Reno
-        self.cwnd += self.mss
+        raise NotImplementedError(
+            "more_dupacks_received must be implemented by subclasses."
+        )
 
-    def set_before_control(self, current_time, packet_in_flight=0):
-        pass
+    def cwnd_in_segments(self) -> float:
+        """Return the congestion window expressed in number of MSS-sized segments."""
+        return self.cwnd / self.mss
 
+    def min_ssthresh(self) -> float:
+        """The RFC 5681-compliant minimum slow-start threshold (2 MSS)."""
+        return 2 * self.mss
 
-class TCPReno(CongestionControl):
-    """TCP Reno, defined in RFC 2001."""
+    def set_before_control(self, current_time, packet_in_flight: int = 0):
+        """Optional hook for controllers that need per-send context (used by BBR)."""
+        _ = (current_time, packet_in_flight)
+
+
+class LossBasedCongestionControl(CongestionControl):
+    """
+    Implements the shared bookkeeping for Reno-like algorithms that respond to loss.
+
+    The derived classes must provide the congestion-avoidance rule via
+    :meth:`_congestion_avoidance_ack` and may override the loss hooks if additional
+    state (e.g., CUBIC's epoch) needs to be updated.
+    """
+
+    beta: Final[float] = 0.5
+    beta_timeout: Final[float] = 0.5
 
     def ack_received(self, rtt: float = 0, current_time: float = 0):
-        """Actions to be taken when a new ack has been received."""
-        if self.cwnd <= self.ssthresh:
-            # slow start
-            self.cwnd += self.mss
+        """RFC 5681 slow start followed by an algorithm-specific avoidance rule."""
+        if self.cwnd < self.ssthresh:
+            self._slow_start_ack()
         else:
-            # congestion avoidance
-            self.cwnd += self.mss * self.mss / self.cwnd
+            self._congestion_avoidance_ack(rtt, current_time)
+
+    def timer_expired(self, packet=None):
+        """RFC 5681 timeout handling."""
+        prev_cwnd = self.cwnd
+        self.ssthresh = self._ssthresh_after_loss(prev_cwnd, LossEvent.TIMEOUT)
+        self.cwnd = self.mss  # reset to one MSS per RFC 5681 §3.1
+        self._after_timeout(prev_cwnd, packet)
+
+    def dupack_over(self):
+        """Exit fast recovery once the lost data is cumulatively acknowledged."""
+        self.cwnd = self.ssthresh
+        self._after_fast_recovery_exit()
+
+    def consecutive_dupacks_received(self, packet=None):
+        """Standard fast retransmit / fast recovery entry."""
+        prev_cwnd = self.cwnd
+        self.ssthresh = self._ssthresh_after_loss(prev_cwnd, LossEvent.FAST_LOSS)
+        # Per RFC 5681 §3.2, inflate the window by 3 segments to keep the ACK clock.
+        self.cwnd = self.ssthresh + 3 * self.mss
+        self._after_fast_loss(prev_cwnd, packet)
+
+    def more_dupacks_received(self, packet=None):
+        """Additional dupacks add one MSS so we clock out a replacement segment."""
+        self.cwnd += self.mss
+        self._during_fast_recovery(packet)
+
+    def _slow_start_ack(self):
+        self.cwnd += self.mss
+
+    @abstractmethod
+    def _congestion_avoidance_ack(self, rtt: float, current_time: float):
+        """Algorithm-specific congestion avoidance (one cwnd increase per RTT)."""
+
+    def _ssthresh_after_loss(self, prev_cwnd: float, event: LossEvent) -> float:
+        factor = self.beta_timeout if event == LossEvent.TIMEOUT else self.beta
+        target = prev_cwnd * (1 - factor)
+        return max(self.min_ssthresh(), target)
+
+    def _after_fast_loss(self, prev_cwnd: float, packet=None):
+        """Hook for algorithms that maintain extra state on fast loss."""
+        _ = (prev_cwnd, packet)
+
+    def _during_fast_recovery(self, packet=None):
+        """Hook invoked for each extra dupack while in fast recovery."""
+        _ = packet
+
+    def _after_fast_recovery_exit(self):
+        """Hook invoked when fast recovery completes."""
+
+    def _after_timeout(self, prev_cwnd: float, packet=None):
+        """Hook invoked after the timeout logic resets cwnd."""
+        _ = (prev_cwnd, packet)
+
+
+class TCPReno(LossBasedCongestionControl):
+    """TCP Reno as defined in RFC 5681."""
+
+    def _congestion_avoidance_ack(self, rtt: float = 0, current_time: float = 0):
+        """Additively increase cwnd by roughly one MSS per RTT."""
+        del rtt, current_time
+        self.cwnd += (self.mss * self.mss) / self.cwnd

ns/flow/cubic.py

@@ -1,26 +1,30 @@
 """
-The TCP CUBIC congestion control algorithm, used in the Linux kernel since 2.6.19.
+TCP CUBIC congestion control (RFC 8312) as used in Linux since v2.6.19.
 
-Reference:
-
-Sangtae Ha; Injong Rhee; Lisong Xu. "CUBIC: A New TCP-Friendly High-Speed TCP Variant,"
-ACM SIGOPS Operating Systems Review. 42 (5): 64–74, July 2008.
+Reference: Sangtae Ha, Injong Rhee, Lisong Xu. "CUBIC: A New TCP-Friendly
+High-Speed TCP Variant," ACM SIGOPS OSR, 42(5):64–74, 2008.
 """
-from ns.flow.cc import CongestionControl
+from __future__ import annotations
+
+from ns.flow.cc import LossBasedCongestionControl
 
 
-class TCPCubic(CongestionControl):
+class TCPCubic(LossBasedCongestionControl):
     """
-    The TCP CUBIC congestion control algorithm, used in the Linux kernel since 2.6.19.
+    TCP CUBIC congestion control with RFC 8312 compliant window update rules.
 
     Parameters
     ----------
     mss: int
-        the maximum segment size
+        The maximum segment size (bytes).
     cwnd: int
-        the size of the congestion window.
+        The initial congestion window (bytes).
     ssthresh: int
-        the slow start threshold.
+        Initial slow-start threshold (bytes).
+    beta: float
+        Multiplicative decrease factor (default 0.2 per RFC 8312 §4.6).
+    cubic_constant: float
+        C in the cubic function W(t) = C(t-K)^3 + W_max (default 0.4).
     debug: bool
         If True, prints more verbose debug information.
     """
@@ -30,88 +34,107 @@ def __init__(
         mss: int = 512,
         cwnd: int = 512,
         ssthresh: int = 65535,
+        beta: float = 0.2,
+        cubic_constant: float = 0.4,
         debug: bool = False,
     ):
         super().__init__(mss, cwnd, ssthresh, debug)
-        self.W_last_max: float = 0
-        self.epoch_start = 0
-        self.origin_point = 0
-        self.d_min: float = 0
-        self.W_tcp = 0
-        self.K = 0
-        self.ack_cnt = 0
+        self.beta = beta
+        self.beta_timeout = beta  # RFC 8312 uses the same decrease on RTO.
+        self.cubic_c = cubic_constant
+
+        # Internal state is tracked in packets (segments) to mirror RFC notation.
+        self.W_last_max: float = 0.0
+        self.epoch_start = 0.0
+        self.origin_point = 0.0
+        self.d_min: float = 0.0
+        self.W_tcp = self.cwnd_in_segments()
+        self.K = 0.0
+        self.ack_cnt = 0.0
         self.tcp_friendliness = True
         self.fast_convergence = True
-        self.beta = 0.2
-        self.C = 0.4
-        self.cwnd_cnt = 0
-        self.cnt = 0
+        self.cwnd_cnt = 0.0
+        self.cnt = float("inf")
 
     def __repr__(self):
         return f"cwnd: {self.cwnd}, ssthresh: {self.ssthresh}"
 
-    def cubic_reset(self):
-        """Resetting the states in CUBIC."""
-        self.W_last_max = 0
-        self.epoch_start = 0
-        self.origin_point = 0
-        self.d_min = 0
-        self.W_tcp = 0
-        self.K = 0
-        self.ack_cnt = 0
+    def ack_received(self, rtt: float = 0, current_time: float = 0):
+        """Record the minimum RTT and defer to LossBased's slow start logic."""
+        if rtt > 0:
+            self.d_min = rtt if self.d_min == 0 else min(self.d_min, rtt)
+        super().ack_received(rtt, current_time)
+
+    def _congestion_avoidance_ack(self, rtt: float, current_time: float):
+        # Without an RTT sample we fall back to Reno-style additive increase.
+        if self.d_min <= 0:
+            self.cwnd += (self.mss * self.mss) / self.cwnd
+            return
 
-    def cubic_update(self, current_time):
-        """Updating CUBIC parameters upon the arrival of a new ack."""
+        self.cubic_update(current_time)
+        ack_threshold = max(1.0, self.cnt)
+        self.cwnd_cnt += 1
+        if self.cwnd_cnt >= ack_threshold:
+            self.cwnd += self.mss
+            self.cwnd_cnt = 0
+
+    def cubic_update(self, current_time: float):
+        """Update the cubic window target (RFC 8312 §4.1)."""
+        cwnd_packets = self.cwnd_in_segments()
         self.ack_cnt += 1
+
         if self.epoch_start <= 0:
             self.epoch_start = current_time
-            if self.cwnd < self.W_last_max:
-                self.K = ((self.W_last_max - self.cwnd) / self.C) ** (1.0 / 3)
-            else:
-                self.K = 0
-                self.origin_point = self.cwnd
             self.ack_cnt = 1
-            self.W_tcp = self.cwnd
+            if cwnd_packets < self.W_last_max:
+                self.K = ((self.W_last_max - cwnd_packets) / self.cubic_c) ** (1.0 / 3.0)
+                self.origin_point = self.W_last_max
+            else:
+                self.K = 0.0
+                self.W_last_max = cwnd_packets
+                self.origin_point = cwnd_packets
+            self.W_tcp = cwnd_packets
+
         t = current_time + self.d_min - self.epoch_start
-        target = self.origin_point + self.C * (t - self.K) ** 3
-        if target > self.cwnd:
-            self.cnt = self.cwnd / (target - self.cwnd)
+        target = self.origin_point + self.cubic_c * (t - self.K) ** 3
+        if target > cwnd_packets:
+            self.cnt = cwnd_packets / max(target - cwnd_packets, 1e-6)
         else:
-            self.cnt = 100 * self.cwnd
+            self.cnt = 100.0 * max(1.0, cwnd_packets)
+
         if self.tcp_friendliness:
-            self.cubic_tcp_friendliness()
+            self.cubic_tcp_friendliness(cwnd_packets)
 
-    def cubic_tcp_friendliness(self):
-        """CUBIC actions in TCP mode."""
-        self.W_tcp += 3 * self.beta / (2 - self.beta) * (self.ack_cnt / self.cwnd)
+    def cubic_tcp_friendliness(self, cwnd_packets: float):
+        """TCP-friendly mode keeps pace with Reno while probing above W_max."""
+        if cwnd_packets <= 0:
+            return
+        self.W_tcp += 3 * self.beta / (2 - self.beta) * (self.ack_cnt / cwnd_packets)
         self.ack_cnt = 0
-        if self.W_tcp > self.cwnd:
-            max_cnt = self.cwnd / (self.W_tcp - self.cwnd)
-            if self.cnt > max_cnt:
-                self.cnt = max_cnt
+        if self.W_tcp > cwnd_packets:
+            max_cnt = cwnd_packets / max(self.W_tcp - cwnd_packets, 1e-6)
+            self.cnt = min(self.cnt, max_cnt)
 
-    def timer_expired(self, packet=None):
-        """Actions to be taken when a timer expired."""
-        # setting the congestion window to 1 segment
-        self.cwnd = self.mss
-        self.cubic_reset()
+    def _reset_epoch(self):
+        """Drop epoch-specific state so the next ACK starts a fresh cubic phase."""
+        self.epoch_start = 0.0
+        self.K = 0.0
+        self.ack_cnt = 0.0
+        self.cnt = float("inf")
+        self.cwnd_cnt = 0.0
+        self.W_tcp = self.cwnd_in_segments()
 
-    def ack_received(self, rtt: float = 0, current_time: float = 0):
-        """Actions to be taken when a new ack has been received."""
-        if self.d_min > 0:
-            self.d_min = min(self.d_min, rtt)
+    def _after_fast_loss(self, prev_cwnd: float, packet=None):
+        """RFC 8312 §4.6 fast convergence bookkeeping."""
+        cwnd_packets = prev_cwnd / self.mss
+        if self.fast_convergence and cwnd_packets < self.W_last_max:
+            self.W_last_max = cwnd_packets * (1 + self.beta) / 2.0
         else:
-            self.d_min = rtt
+            self.W_last_max = cwnd_packets
+        self._reset_epoch()
 
-        if self.cwnd <= self.ssthresh:
-            # slow start
-            self.cwnd += self.mss
-        else:
-            # congestion avoidance
-            self.cubic_update(current_time)
-
-            if self.cwnd_cnt > self.cnt:
-                self.cwnd += self.mss
-                self.cwnd_cnt = 0
-            else:
-                self.cwnd_cnt += 1
+    def _after_timeout(self, prev_cwnd: float, packet=None):
+        """Timeouts force a new epoch but keep the most recent W_max."""
+        del packet
+        self.W_last_max = prev_cwnd / self.mss
+        self._reset_epoch()

tests/conftest.py

@@ -0,0 +1,9 @@
+"""Ensure the project root is importable when tests run from a built wheel."""
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))