Skip to content

Commit 162d9e5

Browse files
committed
Use Reflection for accessing truly private executor method and
use the listener bus to know when receivers have registered (`onStart` is called before receivers have registered, leading to flaky behavior).
1 parent 210f495 commit 162d9e5

File tree

4 files changed

+27
-9
lines changed

4 files changed

+27
-9
lines changed

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

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,6 @@
1717

1818
package org.apache.spark.streaming.receiver
1919

20-
import java.util.concurrent.atomic.AtomicInteger
21-
2220
import com.google.common.util.concurrent.{RateLimiter => GuavaRateLimiter}
2321

2422
import org.apache.spark.{Logging, SparkConf}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -271,7 +271,7 @@ abstract class Receiver[T](val storageLevel: StorageLevel) extends Serializable
271271
}
272272

273273
/** Get the attached executor. */
274-
private[streaming] def executor: ReceiverSupervisor = {
274+
private def executor: ReceiverSupervisor = {
275275
assert(executor_ != null, "Executor has not been attached to this receiver")
276276
executor_
277277
}

streaming/src/main/scala/org/apache/spark/streaming/scheduler/ReceiverTracker.scala

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -182,8 +182,9 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
182182

183183
/** Update a receiver's maximum ingestion rate */
184184
def sendRateUpdate(streamUID: Int, newRate: Long): Unit = {
185-
for (info <- receiverInfo.get(streamUID); eP <- Option(info.endpoint))
185+
for (info <- receiverInfo.get(streamUID); eP <- Option(info.endpoint)) {
186186
eP.send(UpdateRateLimit(newRate))
187+
}
187188
}
188189

189190
/** Add new blocks for the given stream */

streaming/src/test/scala/org/apache/spark/streaming/scheduler/ReceiverTrackerSuite.scala

Lines changed: 24 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -80,12 +80,27 @@ class ReceiverTrackerSuite extends TestSuiteBase {
8080
}
8181

8282
test("Receiver tracker - propagates rate limit") {
83+
object streamingListener extends StreamingListener {
84+
@volatile
85+
var started = false
86+
87+
override def onReceiverStarted(receiverStarted: StreamingListenerReceiverStarted): Unit = {
88+
started = true
89+
}
90+
}
91+
92+
ssc.addStreamingListener(streamingListener)
93+
ssc.scheduler.listenerBus.start(ssc.sc)
94+
8395
val newRateLimit = 100L
8496
val ids = new TestReceiverInputDStream(ssc)
8597
val tracker = new ReceiverTracker(ssc)
8698
tracker.start()
99+
100+
// we wait until the Receiver has registered with the tracker,
101+
// otherwise our rate update is lost
87102
eventually(timeout(5 seconds)) {
88-
assert(TestDummyReceiver.started)
103+
assert(streamingListener.started)
89104
}
90105
tracker.sendRateUpdate(ids.id, newRateLimit)
91106
// this is an async message, we need to wait a bit for it to be processed
@@ -102,7 +117,14 @@ private class TestReceiverInputDStream(@transient ssc_ : StreamingContext)
102117
override def getReceiver(): DummyReceiver = TestDummyReceiver
103118

104119
def getCurrentRateLimit: Option[Long] = {
105-
TestDummyReceiver.executor.getCurrentRateLimit
120+
invokeExecutorMethod.getCurrentRateLimit
121+
}
122+
123+
private def invokeExecutorMethod: ReceiverSupervisor = {
124+
val c = classOf[Receiver[_]]
125+
val ex = c.getDeclaredMethod("executor")
126+
ex.setAccessible(true)
127+
ex.invoke(TestDummyReceiver).asInstanceOf[ReceiverSupervisor]
106128
}
107129
}
108130

@@ -118,10 +140,7 @@ private object TestDummyReceiver extends DummyReceiver
118140
private class DummyReceiver(host: Option[String] = None)
119141
extends Receiver[Int](StorageLevel.MEMORY_ONLY) {
120142

121-
var started = false
122-
123143
def onStart() {
124-
started = true
125144
}
126145

127146
def onStop() {

0 commit comments

Comments
 (0)