diff --git a/runtime/src/main/scala/akka/grpc/internal/RequestBuilderImpl.scala b/runtime/src/main/scala/akka/grpc/internal/RequestBuilderImpl.scala index 02418fea5..ea4f36b63 100644 --- a/runtime/src/main/scala/akka/grpc/internal/RequestBuilderImpl.scala +++ b/runtime/src/main/scala/akka/grpc/internal/RequestBuilderImpl.scala @@ -4,11 +4,11 @@ package akka.grpc.internal -import java.util.concurrent.CompletionStage - +import java.util.concurrent.{ CompletionStage, TimeUnit } import akka.NotUsed import akka.annotation.{ InternalApi, InternalStableApi } import akka.dispatch.ExecutionContexts +import akka.grpc.scaladsl.SingleResponseRequestBuilder import akka.grpc.{ GrpcClientSettings, GrpcResponseMetadata, GrpcServiceException, GrpcSingleResponse } import akka.stream.{ Graph, Materializer, SourceShape } import akka.stream.javadsl.{ Source => JavaSource } @@ -17,6 +17,7 @@ import akka.util.ByteString import io.grpc._ import scala.compat.java8.FutureConverters._ +import scala.concurrent.duration.Duration import scala.concurrent.{ ExecutionContext, Future } /** @@ -52,6 +53,14 @@ final class ScalaUnaryRequestBuilder[I, O]( override def withHeaders(headers: MetadataImpl): ScalaUnaryRequestBuilder[I, O] = new ScalaUnaryRequestBuilder[I, O](descriptor, channel, defaultOptions, settings, headers) + + override def setDeadline(deadline: Duration): SingleResponseRequestBuilder[I, O] = + new ScalaUnaryRequestBuilder[I, O]( + descriptor, + channel, + defaultOptions.withDeadlineAfter(deadline.toMillis, TimeUnit.MILLISECONDS), + settings, + headers) } /** @@ -152,6 +161,14 @@ final class ScalaClientStreamingRequestBuilder[I, O]( override def withHeaders(headers: MetadataImpl): ScalaClientStreamingRequestBuilder[I, O] = new ScalaClientStreamingRequestBuilder[I, O](descriptor, channel, defaultOptions, settings, headers) + + override def setDeadline(deadline: Duration): ScalaClientStreamingRequestBuilder[I, O] = + new ScalaClientStreamingRequestBuilder[I, O]( + descriptor, + channel, + defaultOptions.withDeadlineAfter(deadline.toMillis, TimeUnit.MILLISECONDS), + settings, + headers) } /** diff --git a/runtime/src/main/scala/akka/grpc/scaladsl/RequestBuilder.scala b/runtime/src/main/scala/akka/grpc/scaladsl/RequestBuilder.scala index 22debc4dc..b060b86ca 100644 --- a/runtime/src/main/scala/akka/grpc/scaladsl/RequestBuilder.scala +++ b/runtime/src/main/scala/akka/grpc/scaladsl/RequestBuilder.scala @@ -11,6 +11,7 @@ import akka.stream.scaladsl.Source import akka.util.ByteString import scala.concurrent.Future +import scala.concurrent.duration.Duration /** * Request builder for requests providing per call specific metadata capabilities in @@ -48,6 +49,12 @@ trait SingleResponseRequestBuilder[Req, Res] { * Invoke the gRPC method with the additional metadata added and provide access to response metadata */ def invokeWithMetadata(request: Req): Future[GrpcSingleResponse[Res]] + + /** + * Set the deadline for this call + * @return A new request builder, that will use the supplied deadline when invoked + */ + def setDeadline(deadline: Duration): SingleResponseRequestBuilder[Req, Res] } /**