Making RNG more efficient for Bernoulli sampling usage in congestion

Author: huitseeker

core/src/main/scala/org/apache/spark/util/random/RandomSampler.scala

@@ -143,16 +143,16 @@ class BernoulliCellSampler[T](lb: Double, ub: Double, complement: Boolean = fals
  * @tparam T item type
  */
 @DeveloperApi
-class BernoulliSampler[T: ClassTag](fraction: Double) extends RandomSampler[T, T] {
+class BernoulliSampler[T: ClassTag](fraction: Double,
+                                    private val rng: Random = RandomSampler.newDefaultRNG)
+  extends RandomSampler[T, T] {
 
   /** epsilon slop to avoid failure from floating point jitter */
   require(
     fraction >= (0.0 - RandomSampler.roundingEpsilon)
       && fraction <= (1.0 + RandomSampler.roundingEpsilon),
     s"Sampling fraction ($fraction) must be on interval [0, 1]")
 
-  private val rng: Random = RandomSampler.newDefaultRNG
-
   override def setSeed(seed: Long): Unit = rng.setSeed(seed)
 
   override def sample(items: Iterator[T]): Iterator[T] = {

streaming/src/main/scala/org/apache/spark/streaming/receiver/CongestionStrategyImpl.scala

@@ -20,7 +20,7 @@ package org.apache.spark.streaming.receiver
 import scala.collection.mutable.ArrayBuffer
 import java.util.concurrent.atomic.AtomicInteger
 import java.util.Random
-import org.apache.spark.util.random.BernoulliSampler
+import org.apache.spark.util.random.{RandomSampler, BernoulliSampler}
 
 /**
  * This class provides a congestion strategy that ignores
@@ -79,6 +79,8 @@ class DropCongestionStrategy extends CongestionStrategy {
 
 class SamplingCongestionStrategy extends CongestionStrategy {
 
+  private val rng = Random.newDefaultRNG
+
   private val latestBound = new AtomicInteger(-1)
 
   override def onBlockBoundUpdate(bound: Int): Unit = latestBound.set(bound)
@@ -89,7 +91,7 @@ class SamplingCongestionStrategy extends CongestionStrategy {
     val f = bound.toDouble / currentBuffer.size
     val samplees = currentBuffer.to
 
-    val sampled = new BernoulliSampler(f).sample(samplees.toIterator)
+    val sampled = new BernoulliSampler(f, rng).sample(samplees.toIterator)
 
     currentBuffer.clear()
     currentBuffer ++= sampled