1 /***********************************************************************
2  * Copyright (c) 2013-2024 Commonwealth Computer Research, Inc.
3  * All rights reserved. This program and the accompanying materials
4  * are made available under the terms of the Apache License, Version 2.0
5  * which accompanies this distribution and is available at
6  * http://www.opensource.org/licenses/apache2.0.php.
7  ***********************************************************************/
8 
9 package org.apache.spark.geomesa
10 
11 import com.typesafe.scalalogging.LazyLogging
12 import org.apache.spark.rpc.{RpcCallContext, RpcEndpoint, RpcEnv, RpcTimeout}
13 import org.apache.spark.util.RpcUtils
14 import org.apache.spark.{SparkContext, SparkEnv}
15 import org.geotools.api.feature.simple.SimpleFeatureType
16 import org.locationtech.geomesa.spark.GeoMesaSparkKryoRegistrator
17 import org.locationtech.geomesa.utils.geotools.SimpleFeatureTypes._
18 
19 import java.io.Serializable
20 import scala.reflect.ClassTag
21 import scala.util.{Failure, Success, Try}
22 
23 
24 object GeoMesaSparkKryoRegistratorEndpoint extends LazyLogging {
25 
26   val EnablePropertyKey = "spark.geomesa.kryo.rpc.enable"
27 
28   val EndpointName = "kryo-schema"
29 
30   private lazy val Timeout = RpcUtils.askRpcTimeout(SparkEnv.get.conf)
31   private lazy val EndpointRef = RpcUtils.makeDriverRef(EndpointName, SparkEnv.get.conf, SparkEnv.get.rpcEnv)
32 
33   lazy val Client: KryoClient = Option(SparkEnv.get)
34     .filterNot(_.executorId == SparkContext.DRIVER_IDENTIFIER)
35     .filter(endpointEnabled)
36     .map(_ => ExecutorKryoClient).getOrElse(NoOpKryoClient)
37 
38   private def endpointEnabled(sparkEnv: SparkEnv) =
39     !sparkEnv.conf.get(EnablePropertyKey, "true").equalsIgnoreCase("false")
40 
41   def init(): Unit = {
42     Option(SparkEnv.get).foreach {
43       sparkEnv =>
44         if (endpointEnabled(sparkEnv)) {
45           sparkEnv.executorId match {
46             case SparkContext.DRIVER_IDENTIFIER =>
47               val rpcEnv = sparkEnv.rpcEnv
48               Try(rpcEnv.setupEndpoint(EndpointName, new KryoEndpoint(rpcEnv))) match {
49                 case Success(ref) =>
50                   logger.info(s"$EndpointName rpc endpoint registered on driver ${ref.address}")
51                 case Failure(e: IllegalArgumentException) =>
52                   logger.debug(s"$EndpointName rpc endpoint registration failed, may have been already registered", e)
53                 case Failure(e: Exception) =>
54                   logger.warn(s"$EndpointName rpc endpoint registration failed", e)
55               }
56             case _ => GeoMesaSparkKryoRegistrator.putTypes(Client.getTypes())
57           }
58         } else {
59           logger.debug(s"$EndpointName rpc endpoint disabled")
60         }
61     }
62   }
63 
64   class KryoEndpoint(val rpcEnv: RpcEnv) extends RpcEndpoint {
65     override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = {
66       case message: KryoMessage[_] =>
67         logger.info(s"$message received via rpc from ${context.senderAddress}")
68         context.reply(message.reply)
69     }
70   }
71 
72   trait KryoMessage[R] {
73     def reply: R
74     def ask(timeout: RpcTimeout = Timeout)(implicit c: ClassTag[R]): R = {
75       logger.info(s"$this sent via rpc to ${EndpointRef.address}")
76       val start = System.nanoTime()
77       val result = timeout.awaitResult[R](EndpointRef.ask[R](this, timeout))
78       val delta = (System.nanoTime() - start) / 1000000L
79       logger.info(s"$this response via rpc, $delta ms")
80       result
81     }
82   }
83 
84   class KryoGetTypeMessage(id: Int) extends KryoMessage[Option[(String, String)]] with Serializable {
85     def reply: Option[(String, String)] = Option(GeoMesaSparkKryoRegistrator.getType(id))
86     override def toString: String = s"getType(id=$id)"
87   }
88 
89   class KryoGetTypesMessage() extends KryoMessage[Seq[(String, String)]] with Serializable {
90     def reply: Seq[(String, String)] = GeoMesaSparkKryoRegistrator.getTypes.map(encodeSchema)
91     override def toString: String = s"getTypes()"
92   }
93 
94   class KryoPutTypeMessage(id: Int, name: String, spec: String) extends KryoMessage[Int] with Serializable {
95     def this(id: Int, sft: SimpleFeatureType) = this(id, sft.getTypeName, encodeType(sft))
96     def reply: Int = GeoMesaSparkKryoRegistrator.putType(createType(name, spec))
97     override def toString: String = s"putType(id=$id, name=$name, spec=...)"
98   }
99 
100   trait KryoClient {
101     def getTypes(): Seq[SimpleFeatureType]
102     def getType(id: Int): Option[SimpleFeatureType]
103     def putType(id: Int, schema: SimpleFeatureType): Int
104   }
105 
106   protected object ExecutorKryoClient extends KryoClient {
107     def getTypes(): Seq[SimpleFeatureType] = new KryoGetTypesMessage().ask().map(decodeSchema)
108     def getType(id: Int): Option[SimpleFeatureType] = new KryoGetTypeMessage(id).ask()
109     def putType(id: Int, sft: SimpleFeatureType): Int = new KryoPutTypeMessage(id, sft).ask()
110   }
111 
112   protected object NoOpKryoClient extends KryoClient {
113     def getTypes(): Seq[SimpleFeatureType] = Seq.empty
114     def getType(id: Int) = None
115     def putType(id: Int, sft: SimpleFeatureType): Int = id
116   }
117 
118   implicit def encodeSchema(t: SimpleFeatureType): (String, String) = (t.getTypeName, encodeType(t))
119   implicit def decodeSchema(t: (String, String)): SimpleFeatureType = createType(t._1, t._2)
120   implicit def optionSchema(t: Option[(String, String)]): Option[SimpleFeatureType] = t.map(decodeSchema)
121 
122 }