use a shadow header x-grpc-timeout-propagate in METRICS/PROPAGATE_ONLY

mode

Notable changes:
- In `METRICS_ONLY`, `PROPAGATE_ONLY` mode, do not set `grpc-timeout`
header on the outgoing request as this causes the grpc client to
actually timeout. This can be surprising to those that are only opting
in for a propagation mode. Instead pass on a shadow header
`x-grpc-timeout-propagate`, and the receiver can decide if they will
continue to propagate or enforce it.
- Choose the minimum of [requestDeadline, okHttpClient.callTimeout (only
if it's non-zero)] when setting deadline on the outgoing request.
Also record which was used in metrics.
- Throw a custom `DeadlineExceededException` in the outbound
interceptor, which if unhandled will bubble up as 504 or grpc-status=4
(DEADLINE_EXCEEDED).
- Metrics:
  - Add more buckets to `deadline_duration_ms` metric distribution
  - Skip metrics for `livenesscheckaction` and `readinesscheckaction`

GitOrigin-RevId: 756b1899a5a8b951de6db63e1a93894fb70785fa

Author: sy-lee

misk/api/misk.api

@@ -2516,6 +2516,13 @@ public final class misk/web/proxy/WebProxyEntryKt {
 	public static synthetic fun WebProxyEntry$default (Ljava/lang/String;Ljava/lang/String;ILjava/lang/Object;)Lmisk/web/proxy/WebProxyEntry;
 }
 
+public final class misk/web/requestdeadlines/DeadlineExceededException : java/lang/RuntimeException {
+	public fun <init> ()V
+	public fun <init> (Ljava/lang/String;)V
+	public fun <init> (Ljava/lang/String;Ljava/lang/Throwable;)V
+	public synthetic fun <init> (Ljava/lang/String;Ljava/lang/Throwable;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
+}
+
 public final class misk/web/requestdeadlines/DeadlinePropagationModule : misk/inject/KAbstractModule {
 	public fun <init> ()V
 }

misk/src/main/kotlin/misk/client/DeadlinePropagationInterceptor.kt

@@ -1,21 +1,24 @@
 package misk.client
 
+import com.squareup.wire.GrpcMethod
 import jakarta.inject.Inject
 import jakarta.inject.Singleton
-import misk.exceptions.GatewayTimeoutException
 import misk.grpc.GrpcTimeoutMarshaller
 import misk.scope.ActionScoped
 import misk.web.RequestDeadlineMode
 import misk.web.RequestDeadlinesConfig
 import misk.web.WebConfig
 import misk.web.interceptors.RequestDeadlineInterceptor
-import misk.web.mediatype.MediaTypes
+import misk.web.interceptors.RequestDeadlineInterceptor.Companion.CUSTOM_GRPC_TIMEOUT_PROPAGATE_HEADER
+import misk.web.requestdeadlines.DeadlineExceededException
 import misk.web.requestdeadlines.RequestDeadline
 import misk.web.requestdeadlines.RequestDeadlineMetrics
+import misk.web.requestdeadlines.RequestDeadlineMetrics.SourceLabel.OKHTTP_TIMEOUT
+import misk.web.requestdeadlines.RequestDeadlineMetrics.SourceLabel.PROPAGATED_DEADLINE
 import okhttp3.Interceptor
 import okhttp3.Request
 import okhttp3.Response
-import java.util.concurrent.TimeUnit
+import java.time.Duration
 
 internal class DeadlinePropagationInterceptor(
   private val clientAction: ClientAction,
@@ -26,107 +29,132 @@ internal class DeadlinePropagationInterceptor(
 
   override fun intercept(chain: Interceptor.Chain): Response {
     val requestDeadline = requestDeadlineActionScope.getIfInScope()
+    val isGrpc = chain.request().tag(GrpcMethod::class.java) != null
 
-    // Handle case when no deadline is in scope
+    // Handle case when no deadline is in scope, like if requests were being made from executor or coroutine threads
+    // where the ActionScope has not been explicitly passed along
     if (requestDeadline?.remaining() == null) {
-      return handleNoDeadlineInScope(chain)
+      return handleNoDeadlineInScope(chain, isGrpc)
     }
 
     // Always check deadline and emit metrics based on mode
     // At this point we know requestDeadline is not null (null case handled above)
     return when (requestDeadlinesConfig.mode) {
-      RequestDeadlineMode.METRICS_ONLY -> handleDisabledMode(requestDeadline, chain)
-      RequestDeadlineMode.PROPAGATE_ONLY -> handlePropagateOnlyMode(requestDeadline, chain)
-      RequestDeadlineMode.ENFORCE_INBOUND -> handlePropagateOnlyMode(requestDeadline, chain)
-      RequestDeadlineMode.ENFORCE_OUTBOUND, RequestDeadlineMode.ENFORCE_ALL -> handleEnforceMode(requestDeadline, chain)
+      RequestDeadlineMode.METRICS_ONLY -> handleDisabledMode(requestDeadline, chain, isGrpc)
+      RequestDeadlineMode.PROPAGATE_ONLY, RequestDeadlineMode.ENFORCE_INBOUND -> handlePropagateOnlyMode(requestDeadline, chain, isGrpc)
+      RequestDeadlineMode.ENFORCE_OUTBOUND, RequestDeadlineMode.ENFORCE_ALL -> handleEnforceMode(requestDeadline, chain, isGrpc)
     }
   }
 
-  private fun handleNoDeadlineInScope(chain: Interceptor.Chain): Response {
-    return when (requestDeadlinesConfig.mode) {
-      RequestDeadlineMode.METRICS_ONLY -> {
-        chain.proceed(chain.request())
-      }
-      else -> {
-        // No RequestDeadline found in ActionScope, so propagate client.readTimeoutMillis as deadline
-        val fallbackDeadlineMs: Long = chain.readTimeoutMillis().toLong()
-        metrics.recordOutboundDeadlinePropagated(clientAction, fallbackDeadlineMs, chain.request())
-        val newRequestBuilder = setRequestDeadlineHeadersOnOutbound(chain.request(), fallbackDeadlineMs)
-        chain.proceed(newRequestBuilder.build())
-      }
+  private fun handleNoDeadlineInScope(chain: Interceptor.Chain, isGrpc: Boolean): Response {
+    metrics.recordNoDeadlineInScope(clientAction, isGrpc)
+
+    // For METRICS_ONLY mode or when no fallback deadline exists, proceed with original request
+    if (requestDeadlinesConfig.mode == RequestDeadlineMode.METRICS_ONLY) {
+      return chain.proceed(chain.request())
     }
+
+    val okhttpClientFallbackDeadline = maybeOkHttpClientCallTimeout(chain)
+      ?: return chain.proceed(chain.request())
+
+    val enforced = requestDeadlinesConfig.mode in setOf(RequestDeadlineMode.ENFORCE_OUTBOUND, RequestDeadlineMode.ENFORCE_ALL)
+    metrics.recordOutboundDeadlinePropagated(clientAction, okhttpClientFallbackDeadline, isGrpc, OKHTTP_TIMEOUT)
+    val newRequestBuilder = setRequestDeadlineHeadersOnOutbound(chain.request(), okhttpClientFallbackDeadline, isGrpc, enforced)
+    return chain.proceed(newRequestBuilder.build())
   }
 
-  private fun handleDisabledMode(requestDeadline: RequestDeadline, chain: Interceptor.Chain): Response {
+  private fun handleDisabledMode(requestDeadline: RequestDeadline, chain: Interceptor.Chain, isGrpc: Boolean): Response {
     // Emit metrics, but do not propagate deadline headers or enforce
     if (requestDeadline.expired()) {
       metrics.recordOutboundDeadlineExceeded(
-        clientAction, 
-        enforced = false, 
-        chain.request(),
-        requestDeadline.expiredDuration().toMillis()
-      )
+        clientAction,
+        enforced = false,
+        isGrpc,
+        requestDeadline.expiredDuration())
     }
     return chain.proceed(chain.request())
   }
 
-  private fun handlePropagateOnlyMode(requestDeadline: RequestDeadline, chain: Interceptor.Chain): Response {
+  private fun handlePropagateOnlyMode(requestDeadline: RequestDeadline, chain: Interceptor.Chain, isGrpc: Boolean): Response {
     // Always emit metrics, propagate deadline headers, but do not enforce
+    val enforced = false
     if (requestDeadline.expired()) {
       metrics.recordOutboundDeadlineExceeded(
         clientAction,
-        enforced = false,
-        chain.request(),
-        requestDeadline.expiredDuration().toMillis()
+        enforced,
+        isGrpc,
+        requestDeadline.expiredDuration()
       )
       // Deadline has expired, but config mode specifies it cannot be enforced. For this special case, omit deadline
       // headers, otherwise it will be 0 or a negative number. Let the downstream fallback to a default in its server
       // interceptor.
       return chain.proceed(chain.request())
     } else {
-      val remainingMs = requestDeadline.remaining()!!.toMillis()
-      metrics.recordOutboundDeadlinePropagated(clientAction, remainingMs, chain.request())
-      val requestBuilder = setRequestDeadlineHeadersOnOutbound(chain.request(), remainingMs)
+      val (deadline, source) = determineEffectiveDeadline(requestDeadline, chain)
+      metrics.recordOutboundDeadlinePropagated(clientAction, deadline, isGrpc, source)
+      val requestBuilder = setRequestDeadlineHeadersOnOutbound(chain.request(), deadline, isGrpc, enforced)
       return chain.proceed(requestBuilder.build())
     }
   }
 
-  private fun handleEnforceMode(requestDeadline: RequestDeadline, chain: Interceptor.Chain): Response {
+  private fun handleEnforceMode(requestDeadline: RequestDeadline, chain: Interceptor.Chain, isGrpc: Boolean): Response {
+    val enforced = true
     if (requestDeadline.expired()) {
       metrics.recordOutboundDeadlineExceeded(
         clientAction,
-        enforced = true,
-        chain.request(),
-        requestDeadline.expiredDuration().toMillis()
+        enforced,
+        isGrpc,
+        requestDeadline.expiredDuration()
       )
-      throw GatewayTimeoutException(
+      throw DeadlineExceededException(
         "Deadline already expired, not initiating outbound call to ${chain.request().url}"
       )
     } else {
-      val remainingMs = requestDeadline.remaining()!!.toMillis()
-      metrics.recordOutboundDeadlinePropagated(clientAction, remainingMs, chain.request())
-      val requestBuilder = setRequestDeadlineHeadersOnOutbound(chain.request(), remainingMs)
+      val (deadline, source) = determineEffectiveDeadline(requestDeadline, chain)
+      metrics.recordOutboundDeadlinePropagated(clientAction, deadline, isGrpc, source)
+      val requestBuilder = setRequestDeadlineHeadersOnOutbound(chain.request(), deadline, isGrpc, enforced)
       return chain.proceed(requestBuilder.build())
     }
   }
 
-  private fun setRequestDeadlineHeadersOnOutbound(request: Request, deadlineMs: Long): Request.Builder {
-    val builder = request.newBuilder()
+  /**
+   * Determines the effective deadline by taking the minimum of the request deadline and OkHttp client timeout.
+   * For OkHttpClient timeout, prefer callTimeout if it exists and > 0.
+   * @return Pair of (deadline, source) where deadline is a Duration and source indicates which timeout was used
+   */
+  private fun determineEffectiveDeadline(requestDeadline: RequestDeadline, chain: Interceptor.Chain): Pair<Duration, String> {
+    val propagatedDeadline = requestDeadline.remaining()!!
+    val okHttpClientTimeout = maybeOkHttpClientCallTimeout(chain)
+
+    return if (okHttpClientTimeout == null || propagatedDeadline <= okHttpClientTimeout) {
+      propagatedDeadline to PROPAGATED_DEADLINE
+    } else {
+      okHttpClientTimeout to OKHTTP_TIMEOUT
+    }
+  }
 
-    // nb: Content-Type header not available yet at this point to distinguish http from grpc, so use "te"
-    val isGrpcRequest = request.header("te")?.equals("trailers") == true
-
-    if (isGrpcRequest) {
-      // gRPC request - only add gRPC timeout header
-      if (request.headers.get(GrpcTimeoutMarshaller.TIMEOUT_KEY).isNullOrEmpty()) {
-        builder.header(
-          GrpcTimeoutMarshaller.TIMEOUT_KEY,
-          GrpcTimeoutMarshaller.toAsciiString(TimeUnit.MILLISECONDS.toNanos(deadlineMs)),
-        )
-      }
+  private fun maybeOkHttpClientCallTimeout(chain: Interceptor.Chain): Duration? {
+    val callTimeoutNanos = chain.call().timeout().timeoutNanos()
+    return Duration.ofNanos(callTimeoutNanos).takeIf { callTimeoutNanos != 0L }
+  }
+
+  private fun setRequestDeadlineHeadersOnOutbound(
+    request: Request,
+    deadline: Duration,
+    isGrpc: Boolean,
+    enforced: Boolean
+  ): Request.Builder {
+    val builder = request.newBuilder()
+    if (isGrpc) {
+      // gRPC request - use real header `grpc-timeout` for enforcing modes, shadow header for non-enforcing modes
+      val grpcTimeoutHeader = if (enforced) GrpcTimeoutMarshaller.TIMEOUT_KEY else CUSTOM_GRPC_TIMEOUT_PROPAGATE_HEADER
+      builder.header(
+        grpcTimeoutHeader,
+        GrpcTimeoutMarshaller.toAsciiString(deadline.toNanos()),
+      )
     } else {
-      // HTTP request - only add HTTP deadline header
-      builder.header(RequestDeadlineInterceptor.HTTP_HEADER_ENVOY_DEADLINE, deadlineMs.toString())
+      // HTTP request - only add HTTP deadline header using ISO8601 duration format
+      builder.header(RequestDeadlineInterceptor.HTTP_HEADER_X_REQUEST_DEADLINE, deadline.toString())
     }
 
     return builder

misk/src/main/kotlin/misk/web/MiskWebModule.kt

@@ -41,6 +41,7 @@ import misk.web.actions.StatusAction
 import misk.web.concurrencylimits.ConcurrencyLimiterFactory
 import misk.web.concurrencylimits.ConcurrencyLimitsModule
 import misk.web.exceptions.ActionExceptionLogLevelConfig
+import misk.web.exceptions.DeadlineExceededExceptionMapper
 import misk.web.exceptions.EofExceptionMapper
 import misk.web.exceptions.ExceptionHandlingInterceptor
 import misk.web.exceptions.ExceptionMapperModule
@@ -53,6 +54,7 @@ import misk.web.extractors.PathParamFeatureBinding
 import misk.web.extractors.QueryParamFeatureBinding
 import misk.web.extractors.RequestBodyException
 import misk.web.extractors.RequestBodyFeatureBinding
+import misk.web.requestdeadlines.DeadlineExceededException
 import misk.web.extractors.RequestCookieFeatureBinding
 import misk.web.extractors.RequestCookiesFeatureBinding
 import misk.web.extractors.RequestHeaderFeatureBinding
@@ -266,6 +268,7 @@ class MiskWebModule @JvmOverloads constructor(
     install(ExceptionMapperModule.create<IOException, IOExceptionMapper>())
     install(ExceptionMapperModule.create<EofException, EofExceptionMapper>())
     install(ExceptionMapperModule.create<RequestBodyException, RequestBodyExceptionMapper>())
+    install(ExceptionMapperModule.create<DeadlineExceededException, DeadlineExceededExceptionMapper>())
 
     // Register built-in feature bindings.
     multibind<FeatureBinding.Factory>().toInstance(PathParamFeatureBinding.Factory)

misk/src/main/kotlin/misk/web/exceptions/DeadlineExceededExceptionMapper.kt

@@ -0,0 +1,33 @@
+package misk.web.exceptions
+
+import com.squareup.wire.GrpcStatus.Companion.DEADLINE_EXCEEDED
+import jakarta.inject.Inject
+import misk.client.HTTP_GATEWAY_TIMEOUT
+import misk.web.Response
+import misk.web.mediatype.MediaTypes
+import misk.web.requestdeadlines.DeadlineExceededException
+import misk.web.toResponseBody
+import okhttp3.Headers.Companion.headersOf
+import org.slf4j.event.Level
+
+/**
+ * Maps [DeadlineExceededException] to HTTP 504 Gateway Timeout responses.
+ * This indicates the server was acting as a gateway and didn't receive a timely response.
+ */
+internal class DeadlineExceededExceptionMapper @Inject internal constructor() : ExceptionMapper<DeadlineExceededException> {
+  override fun loggingLevel(th: DeadlineExceededException) = Level.WARN
+
+  override fun toResponse(th: DeadlineExceededException) = DEADLINE_EXCEEDED_RESPONSE
+
+  override fun toGrpcResponse(th: DeadlineExceededException): GrpcErrorResponse? {
+    return GrpcErrorResponse(DEADLINE_EXCEEDED, th.message)
+  }
+
+  companion object {
+    val DEADLINE_EXCEEDED_RESPONSE = Response(
+      "deadline exceeded".toResponseBody(),
+      headersOf("Content-Type", MediaTypes.TEXT_PLAIN_UTF8),
+      HTTP_GATEWAY_TIMEOUT
+    )
+  }
+}

misk/src/main/kotlin/misk/web/interceptors/RequestDeadlineInterceptor.kt

@@ -102,18 +102,24 @@ internal class RequestDeadlineInterceptor private constructor(
 
     return when (action.dispatchMechanism) {
       DispatchMechanism.GRPC -> {
-        val grpcTimeoutString = httpCall.requestHeaders[GrpcTimeoutMarshaller.TIMEOUT_KEY]
-        val grpcTimeoutNanos = grpcTimeoutString?.let { GrpcTimeoutMarshaller.parseAsciiString(it) }?.takeIf { it > 0 }
+        val grpcTimeoutValue = httpCall.requestHeaders[GrpcTimeoutMarshaller.TIMEOUT_KEY]
+        val grpcTimeoutNanos = grpcTimeoutValue?.let { GrpcTimeoutMarshaller.parseAsciiString(it) }?.takeIf { it > 0 }
 
-        if (grpcTimeoutNanos != null) {
-          Duration.ofNanos(grpcTimeoutNanos) to SourceLabel.GRPC_TIMEOUT
+        // "shadow" grpc headers used in METRICS_ONLY/PROPAGATE_ONLY mode to prevent GrpcClient actually hanging up on grpc-timeout
+        val shadowGrpcTimeoutValue = httpCall.requestHeaders[CUSTOM_GRPC_TIMEOUT_PROPAGATE_HEADER]
+        val shadowGrpcTimeoutNanos = shadowGrpcTimeoutValue?.let { GrpcTimeoutMarshaller.parseAsciiString(it) }?.takeIf { it > 0 }
+
+        // Prefer actual to shadow value, though only one of them should be set.
+        val effectiveGrpcTimeout = grpcTimeoutNanos ?: shadowGrpcTimeoutNanos
+        if (effectiveGrpcTimeout != null) {
+          Duration.ofNanos(effectiveGrpcTimeout) to SourceLabel.GRPC_TIMEOUT
         } else {
           fallbackTimeout to fallbackSource
         }
       }
       else -> {
         // HTTP timeout logic with source tracking
-        val xRequestDeadlineTimeout = parseXRequestDeadlineHeader(httpCall.requestHeaders[HTTP_HEADER_X_REQUEST_DEADLINE_SECONDS])
+        val xRequestDeadlineTimeout = parseXRequestDeadlineHeader(httpCall.requestHeaders[HTTP_HEADER_X_REQUEST_DEADLINE])
         val envoyTimeoutMs = httpCall.requestHeaders[HTTP_HEADER_ENVOY_DEADLINE]?.toLongOrNull()
 
         val validTimeouts = mutableListOf<Pair<Duration, String>>()
@@ -162,8 +168,12 @@ internal class RequestDeadlineInterceptor private constructor(
   }
 
   companion object {
-    const val HTTP_HEADER_X_REQUEST_DEADLINE_SECONDS = "X-Request-Deadline"
+    const val HTTP_HEADER_X_REQUEST_DEADLINE = "x-request-deadline"
     const val HTTP_HEADER_ENVOY_DEADLINE = "x-envoy-expected-rq-timeout-ms" // milliseconds
+    // custom header used in METRICS_ONLY/PROPAGATE_ONLY mode to propagate grpc-timeout value, to prevent
+    // GrpcClient actually hanging up
+    const val CUSTOM_GRPC_TIMEOUT_PROPAGATE_HEADER = "x-grpc-timeout-propagate"
+
     const val MISK_REQUEST_DEADLINE_HEADER = "Misk-Request-Deadline"
 
     const val DEADLINE_EXCEEDED_MESSAGE = "deadline exceeded: queued for too long"

misk/src/main/kotlin/misk/web/requestdeadlines/DeadlineExceededException.kt

@@ -0,0 +1,10 @@
+package misk.web.requestdeadlines
+
+/**
+ * Exception thrown when a request deadline has been exceeded.
+ * This is used internally by the deadline propagation system to enforce timeouts.
+ */
+class DeadlineExceededException @JvmOverloads constructor(
+  message: String = "deadline exceeded",
+  cause: Throwable? = null
+) : RuntimeException(message, cause)