Weighted Deterministic Aperture

finagle-core/src/main/scala/com/twitter/finagle/loadbalancer/aperture/Ring.scala

-private class Ring(size: Int, rng: Rng) {
+private class WeightedRing(weights: Vector[Int], rng: Rng) {
...

-  /**
-   * Returns the uniform width of any given index on the ring. The returned
-   * value is bounded between (0, 1].
-   */
-  val unitWidth: Double = 1.0 / size

+  /**
+   * Returns the width of any given index on the ring. The returned
+   * value is bounded between (0, 1].
+   */
+  def unitWidth(ix: Int): Double = weights(ix) / weights.sum

+  /**
+   * Returns the arc length from 0 to index position. The returned
+   * value is bounded between (0, 1].
+   */
+  def widthUntil(ix: Int): Double = weights.slice(0, ix).sum

...

   /**
    * Returns the (zero-based) index between [0, `size`) which the
    * position `offset` maps to.
    *
    * @param offset A value between [0, 1.0).
    */
   def index(offset: Double): Int = {
     if (offset < 0 && offset >= 1.0)
       throw new IllegalArgumentException(s"offset must be between [0, 1.0): $offset")

-    math.floor(offset * size).toInt % size
+    var length = 0
+    for (ix <- weights.indices) {
+      length += unitWidth(ix)
+
+      if (length == offset)
+        return (ix + 1) % weights.length
+      else if (length >= offset)
+        return ix
+    }
+
+    0
   }

   /**
    * Returns the ratio of the intersection between `index` and [offset, offset + width).
    */
   def weight(index: Int, offset: Double, width: Double): Double = {
-    if (index >= size)
+    if (index >= weights.length)
       throw new IllegalArgumentException(s"index must be < size: $index")
     if (width < 0 || width > 1.0)
       throw new IllegalArgumentException(s"width must be between [0, 1.0]: $width")

     // In cases where `offset + width` wraps around the ring, we need
     // to scale the range by 1.0 where it overlaps.
     val ab: Double = {
-      val ab0 = index * unitWidth
+      val ab0 = widthUntil(index)
       if (ab0 + 1 < offset + width) ab0 + 1 else ab0
     }
-    val ae: Double = ab + unitWidth
-    intersect(ab, ae, offset, offset + width) / unitWidth
+    val ae: Double = ab + unitWidth(index)
+    intersect(ab, ae, offset, offset + width) / unitWidth(index)
   }

   /**
    * Pick a random index between [0, `size`) where the range of the
    * index intersects with [offset, offset + width).
    *
    * @param width The width of the range. We interpret a width of 0 as the range
    * [offset, offset] and as such return a valid index.
    */
   def pick(offset: Double, width: Double): Int = {
     if (width < 0 || width > 1.0)
       throw new IllegalArgumentException(s"width must be between [0, 1.0]: $width")

     index((offset + (rng.nextDouble() * width)) % 1.0)
   }

   /**
    * Picks a random index between [0, `size`) where the positions for the
    * respective index intersect with [offset, offset + width), so long as
    * the index is not `a` (if the range permits it).
    *
    * @note we expose this outside of `pick2` so that we can avoid a tuple
    * allocation on the hot path.
    */
   def tryPickSecond(a: Int, offset: Double, width: Double): Int = {
     // Element `b` is picked from "piecewise" range we get by subtracting
     // the range of a, i.e.: [offset, ab), [ae, offset + width).
     // In cases where `offset + width` wraps around the ring, we need
     // to scale the range by 1.0 where it overlaps.
     val ab: Double = {
-      val ab0 = (a * unitWidth)
+      val ab0 = widthUntil(a)
       if (ab0 + 1 < offset + width) ab0 + 1 else ab0
     }
-    val ae: Double = ab + unitWidth
+    val ae: Double = ab + unitWidth(a)

     val overlap = intersect(ab, ae, offset, offset + width)
     val rem = width - overlap

     if (rem > 0) {
       // Instead of actually splitting the range into two, we offset
       // any pick that takes place in the second range if there is a
       // possibility that our second choice falls within [ab, ae].
       //
       // Note, special care must be taken to not bias towards ae + overlap, so
       // we treat the entire range greater than it uniformly.
       var pos = offset + (rng.nextDouble() * rem)
       if (pos >= ae - overlap) { pos += overlap }
       index(pos % 1.0)
     } else {
       // The range [offset, offset + width) is equivalent to [ab, ae).
       a
     }
   }

...

weight(index, offset, width) {
  const a = [
    this.widthUntil(index),
    this.widthUntil(index) + this.unitWidth(index),
  ]
  const b = [offset, offset + width]

  const [first, last] = a[0] < b[0] ? [a, b] : [b, a]

  let i = intersect(...first, ...last)
  // If we wraparound check for any intersection after the wraparound,
  //  it can happen at most once.
  if (last[1] > 1) i += intersect(...first, 0, last[1] % 1.0)

  return i / (1 / this.weights.length)
}