1 /***********************************************************************
2  * Copyright (c) 2013-2025 General Atomics Integrated Intelligence, 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  * https://www.apache.org/licenses/LICENSE-2.0
7  ***********************************************************************/
8 
9 package org.locationtech.geomesa.hbase.server.common
10 
11 import com.google.protobuf.{ByteString, RpcCallback, RpcController}
12 import com.typesafe.scalalogging.StrictLogging
13 import org.apache.hadoop.hbase.client.Scan
14 import org.apache.hadoop.hbase.protobuf.ProtobufUtil
15 import org.apache.hadoop.hbase.protobuf.generated.ClientProtos
16 import org.apache.hadoop.hbase.regionserver.RegionScanner
17 import org.apache.hadoop.hbase.shaded.protobuf.ResponseConverter
18 import org.locationtech.geomesa.hbase.proto.GeoMesaProto
19 import org.locationtech.geomesa.hbase.proto.GeoMesaProto.GeoMesaCoprocessorResponse
20 import org.locationtech.geomesa.hbase.rpc.coprocessor.GeoMesaCoprocessor
21 import org.locationtech.geomesa.hbase.server.common.CoprocessorScan.Aggregator
22 import org.locationtech.geomesa.index.iterators.AggregatingScan
23 import org.locationtech.geomesa.index.iterators.AggregatingScan.AggregateCallback
24 import org.locationtech.geomesa.utils.index.ByteArrays
25 import org.locationtech.geomesa.utils.io.WithClose
26 
27 import java.io.{IOException, InterruptedIOException}
28 import java.util.Base64
29 import scala.util.control.NonFatal
30 
31 /**
32  * Common trait for server side coprocessor class
33  */
34 trait CoprocessorScan extends StrictLogging {
35 
36   import scala.collection.JavaConverters._
37 
38   /**
39    * Gets a scanner for the current region
40    *
41    * @param scan scan
42    * @return
43    */
44   protected def getScanner(scan: Scan): RegionScanner
45 
46   /**
47    * Execute an RPC request against this coprocessor
48    *
49    * @param controller controller
50    * @param request request
51    * @param done callback
52    */
53   protected def execute(
54       controller: RpcController,
55       request: GeoMesaProto.GeoMesaCoprocessorRequest,
56       done: RpcCallback[GeoMesaProto.GeoMesaCoprocessorResponse]): Unit = {
57 
58     val results = GeoMesaCoprocessorResponse.newBuilder()
59 
60     if (request.getVersion != CoprocessorScan.AllowableRequestVersion) {
61       // We cannot handle this request.
62       // Immediately return an empty response indicating the highest response version
63       logger.error(
64         s"Got a coprocessor request with version ${request.getVersion}, " +
65             s"but can only handle ${CoprocessorScan.AllowableRequestVersion}.")
66       results.setVersion(CoprocessorScan.AllowableRequestVersion)
67       done.run(results.build)
68     } else {
69       try {
70         val options = GeoMesaCoprocessor.deserializeOptions(request.getOptions.toByteArray)
71         val timeout = options.get(GeoMesaCoprocessor.TimeoutOpt).map(_.toLong)
72         val yieldPartialResults = options.get(GeoMesaCoprocessor.YieldOpt).exists(_.toBoolean)
73         if (!controller.isCanceled && timeout.forall(_ > System.currentTimeMillis())) {
74           val clas = options(GeoMesaCoprocessor.AggregatorClass)
75           val aggregator = Class.forName(clas).newInstance().asInstanceOf[Aggregator]
76           logger.debug(s"Initializing aggregator $aggregator.")
77           aggregator.init(options)
78 
79           val scan = {
80             val bytes = Base64.getDecoder.decode(options(GeoMesaCoprocessor.ScanOpt))
81             ProtobufUtil.toScan(ClientProtos.Scan.parseFrom(bytes))
82           }
83 
84           WithClose(getScanner(scan)) { scanner =>
85             aggregator.setScanner(scanner)
86             aggregator.aggregate(
87               new CoprocessorAggregateCallback(controller, aggregator, results, yieldPartialResults, timeout))
88           }
89         }
90       } catch {
91         case _: InterruptedException | _: InterruptedIOException => // stop processing, but don't return an error to prevent retries
92         case e: IOException => ResponseConverter.setControllerException(controller, e)
93         case NonFatal(e) => ResponseConverter.setControllerException(controller, new IOException(e))
94       }
95 
96       logger.debug(
97         s"Results total size: ${results.getPayloadList.asScala.map(_.size()).sum}" +
98           s"\n\tBatch sizes: ${results.getPayloadList.asScala.map(_.size()).mkString(", ")}")
99 
100       done.run(results.build)
101     }
102   }
103 
104   /**
105    * Coprocessor aggregate callback
106    *
107    * @param controller controller
108    * @param aggregator aggregator instance
109    * @param results result builder
110    * @param timeout timeout
111    */
112   private class CoprocessorAggregateCallback(
113       controller: RpcController,
114       aggregator: Aggregator,
115       results: GeoMesaCoprocessorResponse.Builder,
116       yieldPartialResults: Boolean,
117       timeout: Option[Long]
118     ) extends AggregateCallback {
119 
120     private val start = System.currentTimeMillis()
121 
122     logger.trace(
123       s"Running first batch on aggregator $aggregator" +
124         timeout.map(t => s" with remaining timeout ${t - System.currentTimeMillis()}ms").getOrElse(""))
125 
126     override def batch(bytes: Array[Byte]): Boolean = {
127       results.addPayload(ByteString.copyFrom(bytes))
128       continue()
129     }
130 
131     override def partial(bytes: => Array[Byte]): Boolean = {
132       if (continue()) { true } else {
133         // add the partial results and stop scanning
134         results.addPayload(ByteString.copyFrom(bytes))
135         false
136       }
137     }
138 
139     private def continue(): Boolean = {
140       if (controller.isCanceled) {
141         logger.warn(s"Stopping aggregator $aggregator due to controller being cancelled")
142         false
143       } else if (timeout.exists(_ < System.currentTimeMillis())) {
144         logger.warn(s"Stopping aggregator $aggregator due to timeout of ${timeout.get}ms")
145         false
146       } else if (yieldPartialResults) {
147         val lastScanned = aggregator.getLastScanned
148         logger.trace(
149           s"Stopping aggregator $aggregator at row ${ByteArrays.printable(lastScanned)} and " +
150               "returning intermediate results")
151         // This check covers the HBase Version Aggregator case
152         if (lastScanned != null && !lastScanned.isEmpty) {
153           results.setLastScanned(ByteString.copyFrom(lastScanned))
154         }
155         false
156       } else {
157         logger.trace(
158           s"Running next batch on aggregator $aggregator " +
159             s"with elapsed time ${System.currentTimeMillis() - start}ms" +
160             timeout.map(t => s" and remaining timeout ${t - System.currentTimeMillis()}ms").getOrElse(""))
161         true
162       }
163     }
164   }
165 }
166 
167 object CoprocessorScan {
168   type Aggregator = HBaseAggregator[_ <: AggregatingScan.Result]
169   val AllowableRequestVersion = 1
170 }