1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
|
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.spark.sql.execution.streaming.state
import scala.collection.mutable
import org.apache.spark.SparkEnv
import org.apache.spark.internal.Logging
import org.apache.spark.rpc.{RpcCallContext, RpcEndpointRef, RpcEnv, ThreadSafeRpcEndpoint}
import org.apache.spark.scheduler.ExecutorCacheTaskLocation
import org.apache.spark.util.RpcUtils
/** Trait representing all messages to [[StateStoreCoordinator]] */
private sealed trait StateStoreCoordinatorMessage extends Serializable
/** Classes representing messages */
private case class ReportActiveInstance(storeId: StateStoreId, host: String, executorId: String)
extends StateStoreCoordinatorMessage
private case class VerifyIfInstanceActive(storeId: StateStoreId, executorId: String)
extends StateStoreCoordinatorMessage
private case class GetLocation(storeId: StateStoreId)
extends StateStoreCoordinatorMessage
private case class DeactivateInstances(storeRootLocation: String)
extends StateStoreCoordinatorMessage
private object StopCoordinator
extends StateStoreCoordinatorMessage
/** Helper object used to create reference to [[StateStoreCoordinator]]. */
private[sql] object StateStoreCoordinatorRef extends Logging {
private val endpointName = "StateStoreCoordinator"
/**
* Create a reference to a [[StateStoreCoordinator]]
*/
def forDriver(env: SparkEnv): StateStoreCoordinatorRef = synchronized {
try {
val coordinator = new StateStoreCoordinator(env.rpcEnv)
val coordinatorRef = env.rpcEnv.setupEndpoint(endpointName, coordinator)
logInfo("Registered StateStoreCoordinator endpoint")
new StateStoreCoordinatorRef(coordinatorRef)
} catch {
case e: IllegalArgumentException =>
val rpcEndpointRef = RpcUtils.makeDriverRef(endpointName, env.conf, env.rpcEnv)
logDebug("Retrieved existing StateStoreCoordinator endpoint")
new StateStoreCoordinatorRef(rpcEndpointRef)
}
}
def forExecutor(env: SparkEnv): StateStoreCoordinatorRef = synchronized {
val rpcEndpointRef = RpcUtils.makeDriverRef(endpointName, env.conf, env.rpcEnv)
logDebug("Retrieved existing StateStoreCoordinator endpoint")
new StateStoreCoordinatorRef(rpcEndpointRef)
}
}
/**
* Reference to a [[StateStoreCoordinator]] that can be used to coordinate instances of
* [[StateStore]]s across all the executors, and get their locations for job scheduling.
*/
private[sql] class StateStoreCoordinatorRef private(rpcEndpointRef: RpcEndpointRef) {
private[state] def reportActiveInstance(
storeId: StateStoreId,
host: String,
executorId: String): Unit = {
rpcEndpointRef.send(ReportActiveInstance(storeId, host, executorId))
}
/** Verify whether the given executor has the active instance of a state store */
private[state] def verifyIfInstanceActive(storeId: StateStoreId, executorId: String): Boolean = {
rpcEndpointRef.askWithRetry[Boolean](VerifyIfInstanceActive(storeId, executorId))
}
/** Get the location of the state store */
private[state] def getLocation(storeId: StateStoreId): Option[String] = {
rpcEndpointRef.askWithRetry[Option[String]](GetLocation(storeId))
}
/** Deactivate instances related to a set of operator */
private[state] def deactivateInstances(storeRootLocation: String): Unit = {
rpcEndpointRef.askWithRetry[Boolean](DeactivateInstances(storeRootLocation))
}
private[state] def stop(): Unit = {
rpcEndpointRef.askWithRetry[Boolean](StopCoordinator)
}
}
/**
* Class for coordinating instances of [[StateStore]]s loaded in executors across the cluster,
* and get their locations for job scheduling.
*/
private class StateStoreCoordinator(override val rpcEnv: RpcEnv) extends ThreadSafeRpcEndpoint {
private val instances = new mutable.HashMap[StateStoreId, ExecutorCacheTaskLocation]
override def receive: PartialFunction[Any, Unit] = {
case ReportActiveInstance(id, host, executorId) =>
instances.put(id, ExecutorCacheTaskLocation(host, executorId))
}
override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = {
case VerifyIfInstanceActive(id, execId) =>
val response = instances.get(id) match {
case Some(location) => location.executorId == execId
case None => false
}
context.reply(response)
case GetLocation(id) =>
context.reply(instances.get(id).map(_.toString))
case DeactivateInstances(loc) =>
val storeIdsToRemove =
instances.keys.filter(_.checkpointLocation == loc).toSeq
instances --= storeIdsToRemove
context.reply(true)
case StopCoordinator =>
stop() // Stop before replying to ensure that endpoint name has been deregistered
context.reply(true)
}
}
|
