akka.http.scaladsl.server.ExceptionHandler Scala Examples

package oharastream.ohara.shabondi.sink

import java.time.{Duration => JDuration}
import java.util.concurrent.TimeUnit

import akka.actor.ActorSystem
import akka.http.scaladsl.model.{ContentTypes, HttpEntity, StatusCodes}
import akka.http.scaladsl.server.{ExceptionHandler, Route}
import com.typesafe.scalalogging.Logger
import oharastream.ohara.common.data.Row
import oharastream.ohara.common.util.Releasable
import oharastream.ohara.shabondi.common.{JsonSupport, RouteHandler, ShabondiUtils}
import org.apache.commons.lang3.StringUtils

import scala.collection.mutable.ArrayBuffer
import scala.compat.java8.DurationConverters._
import scala.concurrent.ExecutionContextExecutor
import scala.concurrent.duration.Duration
import spray.json.DefaultJsonProtocol._
import akka.http.scaladsl.marshallers.sprayjson.SprayJsonSupport._

private[shabondi] object SinkRouteHandler {
  def apply(config: SinkConfig)(implicit actorSystem: ActorSystem) =
    new SinkRouteHandler(config)
}

private[shabondi] class SinkRouteHandler(config: SinkConfig)(implicit actorSystem: ActorSystem) extends RouteHandler {
  implicit private val contextExecutor: ExecutionContextExecutor = actorSystem.dispatcher

  private val log              = Logger(classOf[SinkRouteHandler])
  private[sink] val dataGroups = SinkDataGroups(config)

  def scheduleFreeIdleGroups(interval: JDuration, idleTime: JDuration): Unit =
    actorSystem.scheduler.scheduleWithFixedDelay(Duration(1, TimeUnit.SECONDS), interval.toScala) { () =>
      {
        log.trace("scheduled free group, total group: {} ", dataGroups.size)
        dataGroups.freeIdleGroup(idleTime)
      }
    }

  private val exceptionHandler = ExceptionHandler {
    case ex: Throwable =>
      log.error(ex.getMessage, ex)
      complete((StatusCodes.InternalServerError, ex.getMessage))
  }

  private def fullyPollQueue(queue: RowQueue): Seq[Row] = {
    val buffer    = ArrayBuffer.empty[Row]
    var item: Row = queue.poll()
    while (item != null) {
      buffer += item
      item = queue.poll()
    }
    buffer.toSeq
  }

  private def apiUrl = ShabondiUtils.apiUrl

  def route(): Route = handleExceptions(exceptionHandler) {
    path("groups" / Segment) { groupId =>
      get {
        if (StringUtils.isAlphanumeric(groupId)) {
          val group  = dataGroups.createIfAbsent(groupId)
          val result = fullyPollQueue(group.queue).map(row => JsonSupport.toRowData(row))
          complete(result)
        } else {
          val entity =
            HttpEntity(ContentTypes.`text/plain(UTF-8)`, "Illegal group name, only accept alpha and numeric.")
          complete(StatusCodes.NotAcceptable -> entity)
        }
      } ~ {
        complete(StatusCodes.MethodNotAllowed -> s"Unsupported method, please reference: $apiUrl")
      }
    } ~ {
      complete(StatusCodes.NotFound -> s"Please reference: $apiUrl")
    }
  }

  override def close(): Unit = {
    Releasable.close(dataGroups)
  }
} 

// Copyright (c) 2020 Digital Asset (Switzerland) GmbH and/or its affiliates. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

package com.daml.http.json

import akka.http.scaladsl.marshalling.{Marshaller, ToEntityMarshaller}
import akka.http.scaladsl.model.MediaTypes.`application/json`
import akka.http.scaladsl.model.StatusCodes
import akka.http.scaladsl.server.Directives._
import akka.http.scaladsl.server.ExceptionHandler
import akka.http.scaladsl.unmarshalling.{FromEntityUnmarshaller, Unmarshaller}
import com.daml.util.ExceptionOps._
import scalaz.std.stream.unfold
import scalaz.{@@, Tag}
import spray.json._

import scala.concurrent.{ExecutionContext, Future}


object HttpCodec {
  sealed trait JsonApi
  val JsonApi = Tag.of[JsonApi]

  implicit val jsonExceptionHandler: ExceptionHandler = ExceptionHandler {
    case e: DeserializationException =>
      complete(
        (
          StatusCodes.BadRequest,
          ResponseFormats.errorsJsObject(
            StatusCodes.BadRequest,
            s"JSON parser error: ${e.msg}" +: unfoldCauses(e.cause).map(_.description): _*)))
  }

  private[this] def unfoldCauses(t: Throwable): Seq[Throwable] =
    unfold(t)(tnq => Option(tnq) map (tnn => (tnn, tnn.getCause)))

  private[this] val simpleJsValueUnmarshaller =
    Unmarshaller[String, JsValue] { implicit ec: ExecutionContext => value =>
      Future(value.parseJson)
    }

  implicit val jsValueUnmarshaller: FromEntityUnmarshaller[JsValue] =
    (implicitly[FromEntityUnmarshaller[String]] andThen simpleJsValueUnmarshaller
      forContentTypes `application/json`)

  implicit val jsValueMarshaller: ToEntityMarshaller[JsValue] =
    Marshaller.combined((_: JsValue).compactPrint)(Marshaller.stringMarshaller(`application/json`))

  implicit def jsonCodecUnmarshaller[A: RootJsonReader]: FromEntityUnmarshaller[A @@ JsonApi] =
    JsonApi.subst[FromEntityUnmarshaller, A](jsValueUnmarshaller map (_.convertTo[A]))

  implicit def jsonCodecMarshaller[A: RootJsonWriter]: ToEntityMarshaller[A @@ JsonApi] =
    JsonApi.subst[ToEntityMarshaller, A](Marshaller.combined((_: A).toJson))
} 

package io.opencensus.scala.akka.http

import akka.http.scaladsl.model.{HttpHeader, HttpRequest}
import akka.http.scaladsl.server.Directives._
import akka.http.scaladsl.server.{Directive0, Directive1, ExceptionHandler}
import com.typesafe.scalalogging.LazyLogging
import io.opencensus.scala.Tracing
import io.opencensus.scala.akka.http.propagation.AkkaB3FormatPropagation
import io.opencensus.scala.akka.http.trace.EndSpanResponse
import io.opencensus.scala.akka.http.trace.HttpExtractors._
import io.opencensus.scala.http.{HttpAttributes, ServiceAttributes, ServiceData}
import io.opencensus.scala.http.propagation.Propagation
import io.opencensus.trace.{Span, Status}

import scala.util.control.NonFatal

trait TracingDirective extends LazyLogging {

  protected def tracing: Tracing
  protected def propagation: Propagation[HttpHeader, HttpRequest]

  
  def traceRequestForServiceNoSpan(serviceData: ServiceData): Directive0 =
    traceRequest(serviceData).map(_ => ())

  private def traceRequest(serviceData: ServiceData): Directive1[Span] =
    extractRequest.flatMap { req =>
      val span = buildSpan(req, serviceData)
      recordSuccess(span) & recordException(span) & provide(span)
    }

  private def buildSpan(req: HttpRequest, serviceData: ServiceData): Span = {
    val name = req.uri.path.toString()

    val span = propagation
      .extractContext(req)
      .fold(
        { error =>
          logger.debug("Extracting of parent context failed", error)
          tracing.startSpan(name)
        },
        tracing.startSpanWithRemoteParent(name, _)
      )

    ServiceAttributes.setAttributesForService(span, serviceData)
    HttpAttributes.setAttributesForRequest(span, req)
    span
  }

  private def recordSuccess(span: Span) =
    mapResponse(EndSpanResponse.forServer(tracing, _, span))

  private def recordException(span: Span) =
    handleExceptions(ExceptionHandler {
      case NonFatal(ex) =>
        tracing.endSpan(span, Status.INTERNAL)
        throw ex
    })
}

object TracingDirective extends TracingDirective {
  override protected def tracing: Tracing = Tracing
  override protected def propagation: Propagation[HttpHeader, HttpRequest] =
    AkkaB3FormatPropagation
} 

package org.apache.s2graph.http

import akka.http.scaladsl.model.{ContentTypes, HttpEntity, HttpResponse, StatusCodes}
import akka.http.scaladsl.server.Directives._
import akka.http.scaladsl.server.{ExceptionHandler, Route}
import com.fasterxml.jackson.core.JsonParseException
import org.apache.s2graph.core.rest.RequestParser
import org.apache.s2graph.core.storage.MutateResponse
import org.apache.s2graph.core.{GraphElement, S2Graph}
import org.slf4j.LoggerFactory
import play.api.libs.json.{JsValue, Json}

import scala.concurrent.{ExecutionContext, Future}

trait S2GraphMutateRoute extends PlayJsonSupport {

  val s2graph: S2Graph
  val logger = LoggerFactory.getLogger(this.getClass)

  lazy val parser = new RequestParser(s2graph)

  lazy val exceptionHandler = ExceptionHandler {
    case ex: JsonParseException => complete(StatusCodes.BadRequest -> ex.getMessage)
    case ex: java.lang.IllegalArgumentException => complete(StatusCodes.BadRequest -> ex.getMessage)
  }

  lazy val mutateVertex = path("vertex" / Segments) { params =>
    implicit val ec = s2graph.ec

    val (operation, serviceNameOpt, columnNameOpt) = params match {
      case operation :: serviceName :: columnName :: Nil => (operation, Option(serviceName), Option(columnName))
      case operation :: Nil => (operation, None, None)
      case _ => throw new RuntimeException("invalid params")
    }

    entity(as[JsValue]) { payload =>
      val future = vertexMutate(payload, operation, serviceNameOpt, columnNameOpt).map(Json.toJson(_))

      complete(future)
    }
  }

  lazy val mutateEdge = path("edge" / Segment) { operation =>
    implicit val ec = s2graph.ec

    entity(as[JsValue]) { payload =>
      val future = edgeMutate(payload, operation, withWait = true).map(Json.toJson(_))

      complete(future)
    }
  }

  def vertexMutate(jsValue: JsValue,
                   operation: String,
                   serviceNameOpt: Option[String] = None,
                   columnNameOpt: Option[String] = None,
                   withWait: Boolean = true)(implicit ec: ExecutionContext): Future[Seq[Boolean]] = {
    val vertices = parser.toVertices(jsValue, operation, serviceNameOpt, columnNameOpt)

    val verticesToStore = vertices.filterNot(_.isAsync)

    s2graph.mutateVertices(verticesToStore, withWait).map(_.map(_.isSuccess))
  }

  def edgeMutate(elementsWithTsv: Seq[(GraphElement, String)], withWait: Boolean)(implicit ec: ExecutionContext): Future[Seq[Boolean]] = {
    val elementWithIdxs = elementsWithTsv.zipWithIndex
    val (elementSync, elementAsync) = elementWithIdxs.partition { case ((element, tsv), idx) => !element.isAsync }

    val retToSkip = elementAsync.map(_._2 -> MutateResponse.Success)
    val (elementsToStore, _) = elementSync.map(_._1).unzip
    val elementsIdxToStore = elementSync.map(_._2)

    s2graph.mutateElements(elementsToStore, withWait).map { mutateResponses =>
      elementsIdxToStore.zip(mutateResponses) ++ retToSkip
    }.map(_.sortBy(_._1).map(_._2.isSuccess))
  }

  def edgeMutate(jsValue: JsValue, operation: String, withWait: Boolean)(implicit ec: ExecutionContext): Future[Seq[Boolean]] = {
    val edgesWithTsv = parser.parseJsonFormat(jsValue, operation)
    edgeMutate(edgesWithTsv, withWait)
  }

  // expose routes
  lazy val mutateRoute: Route =
    post {
      concat(
        handleExceptions(exceptionHandler) {
          mutateVertex
        },
        handleExceptions(exceptionHandler) {
          mutateEdge
        }
      )
    }

} 

package hydra.ingest.http.mock

import akka.actor.ActorSystem
import akka.http.scaladsl.model.StatusCodes._
import akka.http.scaladsl.server.Directives._
import akka.http.scaladsl.server.{ExceptionHandler, Route}
import hydra.ingest.http.SchemasEndpoint
import io.confluent.kafka.schemaregistry.client.rest.exceptions.RestClientException

import scala.concurrent.{ExecutionContext, Future}

class MockEndpoint(
    implicit system: ActorSystem,
    implicit val e: ExecutionContext
) {

  def throwRestClientException(
      statusCode: Int,
      errorCode: Int,
      errorMessage: String
  ): Future[Any] = {
    throw new RestClientException(errorMessage, statusCode, errorCode)
  }

  val schemaRouteExceptionHandler: ExceptionHandler =
    new SchemasEndpoint().excptHandler

  def route: Route = {
    pathPrefix("throwRestClientException") {
      handleExceptions(schemaRouteExceptionHandler) {
        get {
          parameters('statusCode, 'errorCode, 'errorMessage) {
            (statusCode, errorCode, errorMessage) =>
              pathEndOrSingleSlash {
                onSuccess(
                  throwRestClientException(
                    statusCode.toInt,
                    errorCode.toInt,
                    errorMessage
                  )
                ) { _ => complete(OK) }
              }
          }
        }
      }
    }
  }
} 

package com.shashank.akkahttp.basic.routing

import akka.actor.ActorSystem
import akka.http.scaladsl.model.StatusCodes
import akka.http.scaladsl.Http
import akka.http.scaladsl.model.HttpResponse
import akka.http.scaladsl.server.Directives._
import akka.http.scaladsl.server.ExceptionHandler
import akka.stream.{ActorMaterializer, Materializer}


object Failure {

  def main(args: Array[String]) {

    implicit val sys = ActorSystem("IntroductionToAkkaHttp")
    implicit val mat:Materializer = ActorMaterializer()

    implicit def myExceptionHandler = ExceptionHandler {
      case _: ArithmeticException =>
        complete(HttpResponse(StatusCodes.BadRequest, entity = "Bad numbers, bad result!!!"))
      case e: Throwable => {
        println(e.getMessage)
        println(e.getStackTraceString)
        complete(HttpResponse(StatusCodes.BadRequest, entity = e.getMessage))
      }
    }

    val route =
      path("welcome"){
        get{
          complete {
            "welcome to rest service"
          }
        }
      } ~
      path("demo"){
        get {
          complete {
            100/0
            "welcome to demonstration"
          }
        }
      }

    Http().bindAndHandle(route, "localhost", 8090)
  }

} 

package com.packt.chapter9

import akka.http.scaladsl.model.StatusCodes
import akka.http.scaladsl.server.Directives._
import akka.http.scaladsl.server.ExceptionHandler
import akka.pattern.AskTimeoutException

trait RouteExceptionHandler {

  val routeExceptionHandler = ExceptionHandler {
    case _: ArithmeticException =>
      complete {
        StatusCodes.BadRequest -> "You values are incorrect. Probably b needs to be different from 0"
      }
    case _: AskTimeoutException =>
      complete {
        StatusCodes.ServiceUnavailable -> "Internal actor is not responding within 500 millis"
      }
  }
} 

package com.lonelyplanet.akka.http.extensions

import akka.http.scaladsl.server.ExceptionHandler
import akka.http.scaladsl.server.directives.BasicDirectives._
import com.lonelyplanet.akka.http.extensions.exceptions.ResourceNotFound
import com.lonelyplanet.util.AirbrakeService

trait ExceptionHandlingWithAirbrake extends ExceptionHandling {
  val airbrakeService: AirbrakeService
  val airbrakeComponent: String = ExceptionHandlingWithAirbrake.DefaultAirbrakeComponent

  override implicit def exceptionHandler: ExceptionHandler = ExceptionHandler {
    case e: ResourceNotFound =>
      super.exceptionHandler.apply(e)
    case e =>
      extractUri { uri =>
        airbrakeService.notify(e, uri.toString(), airbrakeComponent)
        super.exceptionHandler.apply(e)
      }
  }
}

object ExceptionHandlingWithAirbrake {
  val DefaultAirbrakeComponent = "api"
} 

package ch.epfl.bluebrain.nexus.storage.routes

import akka.http.scaladsl.model.headers.{`WWW-Authenticate`, HttpChallenges}
import akka.http.scaladsl.server.Directives._
import akka.http.scaladsl.server.{ExceptionHandler, RejectionHandler, Route}
import ch.epfl.bluebrain.nexus.storage.IamIdentitiesClient.Caller
import ch.epfl.bluebrain.nexus.storage.StorageError._
import ch.epfl.bluebrain.nexus.storage.config.AppConfig
import ch.epfl.bluebrain.nexus.storage.config.AppConfig._
import ch.epfl.bluebrain.nexus.storage.routes.AuthDirectives._
import ch.epfl.bluebrain.nexus.storage.routes.PrefixDirectives._
import ch.epfl.bluebrain.nexus.storage.routes.instances._
import ch.epfl.bluebrain.nexus.storage.{AkkaSource, IamIdentitiesClient, Rejection, StorageError, Storages}
import com.typesafe.scalalogging.Logger
import monix.eval.Task

import scala.util.control.NonFatal

object Routes {

  private[this] val logger = Logger[this.type]

  
  def apply(
      storages: Storages[Task, AkkaSource]
  )(implicit config: AppConfig, identities: IamIdentitiesClient[Task]): Route =
    //TODO: Fetch Bearer token and verify identity
    wrap {
      concat(
        AppInfoRoutes(config.description).routes,
        (pathPrefix(config.http.prefix) & extractToken) { implicit token =>
          extractCaller.apply {
            case Caller(config.subject.subjectValue, _) => StorageRoutes(storages).routes
            case _                                      => failWith(AuthenticationFailed)
          }
        }
      )
    }

} 

package com.lonelyplanet.prometheus.directives

import akka.http.scaladsl.server.ExceptionHandler
import akka.http.scaladsl.server.directives.{BasicDirectives, ExecutionDirectives}
import com.lonelyplanet.prometheus.ResponseTimeRecorder

import scala.concurrent.duration
import scala.concurrent.duration.FiniteDuration
import scala.util.control.NonFatal

trait ResponseTimeRecordingDirectives {
  this: ResponseTimeRecorderProvider =>

  
  def recordResponseTime(endpoint: String) = BasicDirectives.extractRequestContext.flatMap { ctx =>
    val requestStartTime = System.nanoTime()
    BasicDirectives.mapResponse { resp =>
      record(endpoint, requestStartTime)
      resp
    } & ExecutionDirectives.handleExceptions {
      responseTimeRecordingExceptionHandler(endpoint, requestStartTime)
    }
  }

  private def responseTimeRecordingExceptionHandler(endpoint: String, requestStartTime: Long) = ExceptionHandler {
    case NonFatal(e) =>
      record(endpoint, requestStartTime)

      // Rethrow the exception to allow proper handling
      // from handlers higher ip in the hierarchy
      throw e
  }

  private def record(endpoint: String, requestStartTime: Long): Unit = {
    val requestEndTime = System.nanoTime()
    val total = new FiniteDuration(requestEndTime - requestStartTime, duration.NANOSECONDS)

    recorder.recordResponseTime(endpoint, total)
  }
}

object ResponseTimeRecordingDirectives {
  def apply(r: ResponseTimeRecorder) = {
    new ResponseTimeRecordingDirectives with ResponseTimeRecorderProvider {
      override def recorder: ResponseTimeRecorder = r
    }
  }
}

trait ResponseTimeRecorderProvider {
  def recorder: ResponseTimeRecorder
} 

package org.squbs.testkit

import akka.actor.{Actor, ActorSystem}
import akka.http.scaladsl.server.directives.PathDirectives._
import akka.http.scaladsl.server.{ExceptionHandler, RejectionHandler, Route}
import akka.http.scaladsl.settings.RoutingSettings
import akka.testkit.TestActorRef
import org.squbs.unicomplex.{RouteDefinition, WithActorContext, WithWebContext}

import scala.reflect.ClassTag


  def apply[T <: RouteDefinition: ClassTag](webContext: String)(implicit system: ActorSystem): Route = {
    val clazz = implicitly[ClassTag[T]].runtimeClass
    implicit val actorContext = TestActorRef[TestRouteActor].underlyingActor.context
    val routeDef =
      WithWebContext(webContext) {
        WithActorContext {
          clazz.asSubclass(classOf[RouteDefinition]).newInstance()
        }
      }

    implicit val routingSettings = RoutingSettings(system.settings.config)
    implicit val rejectionHandler:RejectionHandler = routeDef.rejectionHandler.getOrElse(RejectionHandler.default)
    implicit val exceptionHandler:ExceptionHandler = routeDef.exceptionHandler.orNull

    if (webContext.length > 0) pathPrefix(separateOnSlashes(webContext)) { Route.seal(routeDef.route) }
    else { Route.seal(routeDef.route) }
  }

  private[testkit] class TestRouteActor extends Actor {
    def receive = {
      case _ =>
    }
  }
} 

package org.marvin.executor.api

import akka.http.scaladsl.model._
import akka.http.scaladsl.server.Directives._
import akka.http.scaladsl.server.{ExceptionHandler, MissingQueryParamRejection, RejectionHandler}
import grizzled.slf4j.Logger
import org.marvin.exception.MarvinEExecutorException
import spray.json.DefaultJsonProtocol._

import scala.concurrent.TimeoutException

case class ErrorResponse(errorMessage: String)

object GenericAPIHandlers {

  lazy val logger = Logger[this.type]

  implicit val errorFormatter = jsonFormat1(ErrorResponse)

  val exceptions: ExceptionHandler =
    ExceptionHandler {
      case ex: IllegalArgumentException => {
        logger.debug("Endpoint thrown illegal argument exception.", ex)

        val error = ErrorResponse(errorMessage = ex.getMessage)
        complete(HttpResponse(StatusCodes.BadRequest, entity = toResponseEntityJson(error)))
      }
      case ex: TimeoutException => {
        logger.debug("Endpoint thrown timeout exception", ex)

        val error = ErrorResponse(errorMessage = "The engine was not able to provide a response within the specified timeout.")
        complete(HttpResponse(StatusCodes.InternalServerError, entity = toResponseEntityJson(error)))
      }
      case ex: MarvinEExecutorException => {
        logger.debug("Endpoint thrown Marvin EExecutor Exception", ex)

        val error = ErrorResponse(errorMessage = ex.getMessage())
        complete(HttpResponse(StatusCodes.ServiceUnavailable, entity = toResponseEntityJson(error)))
      }
      case _ => {
        val error = ErrorResponse(errorMessage = "Unexpected error.")
        complete(HttpResponse(StatusCodes.InternalServerError, entity = toResponseEntityJson(error)))
      }
    }

  val rejections: RejectionHandler =
    RejectionHandler.newBuilder().handle {
      case rj: MissingQueryParamRejection => {
        logger.debug("Missing query parameters.")

        val error = ErrorResponse(errorMessage = s"Missing query parameter. [${rj.parameterName}]")
        complete(HttpResponse(StatusCodes.BadRequest, entity = toResponseEntityJson(error)))
      }
    }.result()

  def toResponseEntityJson(error: ErrorResponse): ResponseEntity = {
    HttpEntity(ContentTypes.`application/json`, errorFormatter.write(error).toString())
  }
} 

package sttp.client.akkahttp

import akka.actor.ActorSystem
import akka.event.LoggingAdapter
import akka.http.scaladsl.model.ws.{Message, WebSocketRequest, WebSocketUpgradeResponse}
import akka.http.scaladsl.model.{HttpRequest, HttpResponse}
import akka.http.scaladsl.server.{ExceptionHandler, RejectionHandler, Route, RoutingLog}
import akka.http.scaladsl.settings.{ClientConnectionSettings, ConnectionPoolSettings, ParserSettings, RoutingSettings}
import akka.http.scaladsl.{Http, HttpsConnectionContext}
import akka.stream.Materializer
import akka.stream.scaladsl.Flow

import scala.concurrent.{ExecutionContext, ExecutionContextExecutor, Future}

trait AkkaHttpClient {
  def singleRequest(
      request: HttpRequest,
      settings: ConnectionPoolSettings
  ): Future[HttpResponse]

  def singleWebsocketRequest[WS_RESULT](
      request: WebSocketRequest,
      clientFlow: Flow[Message, Message, WS_RESULT],
      settings: ClientConnectionSettings
  )(implicit ec: ExecutionContext, mat: Materializer): Future[(WebSocketUpgradeResponse, WS_RESULT)]
}

object AkkaHttpClient {
  def default(
      system: ActorSystem,
      connectionContext: Option[HttpsConnectionContext],
      customLog: Option[LoggingAdapter]
  ): AkkaHttpClient =
    new AkkaHttpClient {
      private val http = Http()(system)

      override def singleRequest(
          request: HttpRequest,
          settings: ConnectionPoolSettings
      ): Future[HttpResponse] = {
        http.singleRequest(
          request,
          connectionContext.getOrElse(http.defaultClientHttpsContext),
          settings,
          customLog.getOrElse(system.log)
        )
      }

      override def singleWebsocketRequest[WS_RESULT](
          request: WebSocketRequest,
          clientFlow: Flow[Message, Message, WS_RESULT],
          settings: ClientConnectionSettings
      )(implicit ec: ExecutionContext, mat: Materializer): Future[(WebSocketUpgradeResponse, WS_RESULT)] = {
        val (wsResponse, wsResult) = http.singleWebSocketRequest(
          request,
          clientFlow,
          connectionContext.getOrElse(http.defaultClientHttpsContext),
          None,
          settings,
          customLog.getOrElse(system.log)
        )
        wsResponse.map((_, wsResult))
      }
    }

  def stubFromAsyncHandler(run: HttpRequest => Future[HttpResponse]): AkkaHttpClient =
    new AkkaHttpClient {
      def singleRequest(request: HttpRequest, settings: ConnectionPoolSettings): Future[HttpResponse] =
        run(request)

      override def singleWebsocketRequest[WS_RESULT](
          request: WebSocketRequest,
          clientFlow: Flow[Message, Message, WS_RESULT],
          settings: ClientConnectionSettings
      )(implicit ec: ExecutionContext, mat: Materializer): Future[(WebSocketUpgradeResponse, WS_RESULT)] =
        Future.failed(new RuntimeException("Websockets are not supported"))
    }

  def stubFromRoute(route: Route)(implicit
      routingSettings: RoutingSettings,
      parserSettings: ParserSettings,
      materializer: Materializer,
      routingLog: RoutingLog,
      executionContext: ExecutionContextExecutor = null,
      rejectionHandler: RejectionHandler = RejectionHandler.default,
      exceptionHandler: ExceptionHandler = null
  ): AkkaHttpClient = stubFromAsyncHandler(Route.asyncHandler(route))
} 

package com.wavesplatform.api.http

import akka.http.scaladsl.marshalling.ToResponseMarshallable
import akka.http.scaladsl.server.{Directive1, ExceptionHandler, Route}
import com.google.common.util.concurrent.{ExecutionError, UncheckedExecutionException}
import com.wavesplatform.utils.Schedulers
import monix.execution.Scheduler

import scala.concurrent.ExecutionException

trait TimeLimitedRoute { self: ApiRoute =>
  def limitedScheduler: Scheduler

  def executeLimited[T](f: => T): Directive1[T] = {
    val handler = ExceptionHandler {
      case _: InterruptedException | _: ExecutionException | _: ExecutionError | _: UncheckedExecutionException =>
        complete(ApiError.CustomValidationError("The request took too long to complete"))
    }
    handleExceptions(handler) & onSuccess(Schedulers.executeCatchingInterruptedException(limitedScheduler)(f))
  }

  def completeLimited(f: => ToResponseMarshallable): Route =
    executeLimited(f)(complete(_))
}