From e6c8d02d044081d9945eda83f9bedac985cab11f Mon Sep 17 00:00:00 2001
From: Patrick Koenig <pkoenig10@gmail.com>
Date: Wed, 3 Jan 2024 12:26:54 -0500
Subject: [PATCH] Add reason to propagated QoS errors

---
 .../java/dialogue/serde/ErrorDecoder.java     | 11 +++++---
 .../java/dialogue/serde/ErrorDecoderTest.java | 25 +++++++++++--------
 2 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/dialogue-serde/src/main/java/com/palantir/conjure/java/dialogue/serde/ErrorDecoder.java b/dialogue-serde/src/main/java/com/palantir/conjure/java/dialogue/serde/ErrorDecoder.java
index aaf3cd502b..e0e0612368 100644
--- a/dialogue-serde/src/main/java/com/palantir/conjure/java/dialogue/serde/ErrorDecoder.java
+++ b/dialogue-serde/src/main/java/com/palantir/conjure/java/dialogue/serde/ErrorDecoder.java
@@ -17,11 +17,13 @@
 package com.palantir.conjure.java.dialogue.serde;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.common.io.CharStreams;
 import com.google.common.net.HttpHeaders;
 import com.google.common.primitives.Longs;
 import com.palantir.conjure.java.api.errors.QosException;
+import com.palantir.conjure.java.api.errors.QosReason;
 import com.palantir.conjure.java.api.errors.RemoteException;
 import com.palantir.conjure.java.api.errors.SerializableError;
 import com.palantir.conjure.java.api.errors.UnknownRemoteException;
@@ -57,6 +59,9 @@ public enum ErrorDecoder {
     private static final SafeLogger log = SafeLoggerFactory.get(ErrorDecoder.class);
     private static final ObjectMapper MAPPER = ObjectMappers.newClientObjectMapper();
 
+    @VisibleForTesting
+    static final QosReason QOS_REASON = QosReason.of("client-qos-response");
+
     public boolean isError(Response response) {
         return 300 <= response.code() && response.code() <= 599;
     }
@@ -99,10 +104,10 @@ private RuntimeException decodeInternal(Response response) {
                 return response.getFirstHeader(HttpHeaders.RETRY_AFTER)
                         .map(Longs::tryParse)
                         .map(Duration::ofSeconds)
-                        .map(QosException::throttle)
-                        .orElseGet(QosException::throttle);
+                        .map(duration -> QosException.throttle(QOS_REASON, duration))
+                        .orElseGet(() -> QosException.throttle(QOS_REASON));
             case 503:
-                return QosException.unavailable();
+                return QosException.unavailable(QOS_REASON);
         }
 
         String body;
diff --git a/dialogue-serde/src/test/java/com/palantir/conjure/java/dialogue/serde/ErrorDecoderTest.java b/dialogue-serde/src/test/java/com/palantir/conjure/java/dialogue/serde/ErrorDecoderTest.java
index 249ce8864f..3b1c754434 100644
--- a/dialogue-serde/src/test/java/com/palantir/conjure/java/dialogue/serde/ErrorDecoderTest.java
+++ b/dialogue-serde/src/test/java/com/palantir/conjure/java/dialogue/serde/ErrorDecoderTest.java
@@ -96,7 +96,9 @@ public void testQos503() {
         assertThat(decoder.isError(response)).isTrue();
 
         RuntimeException result = decoder.decode(response);
-        assertThat(result).isInstanceOf(QosException.Unavailable.class);
+        assertThat(result).isInstanceOfSatisfying(QosException.Unavailable.class, exception -> {
+            assertThat(exception.getReason()).isEqualTo(ErrorDecoder.QOS_REASON);
+        });
     }
 
     @Test
@@ -105,9 +107,10 @@ public void testQos429() {
         assertThat(decoder.isError(response)).isTrue();
 
         RuntimeException result = decoder.decode(response);
-        assertThat(result)
-                .isInstanceOfSatisfying(QosException.Throttle.class, exception -> assertThat(exception.getRetryAfter())
-                        .isEmpty());
+        assertThat(result).isInstanceOfSatisfying(QosException.Throttle.class, exception -> {
+            assertThat(exception.getReason()).isEqualTo(ErrorDecoder.QOS_REASON);
+            assertThat(exception.getRetryAfter()).isEmpty();
+        });
     }
 
     @Test
@@ -117,9 +120,10 @@ public void testQos429_retryAfter() {
         assertThat(decoder.isError(response)).isTrue();
 
         RuntimeException result = decoder.decode(response);
-        assertThat(result)
-                .isInstanceOfSatisfying(QosException.Throttle.class, exception -> assertThat(exception.getRetryAfter())
-                        .hasValue(Duration.ofSeconds(3)));
+        assertThat(result).isInstanceOfSatisfying(QosException.Throttle.class, exception -> {
+            assertThat(exception.getReason()).isEqualTo(ErrorDecoder.QOS_REASON);
+            assertThat(exception.getRetryAfter()).hasValue(Duration.ofSeconds(3));
+        });
     }
 
     @Test
@@ -129,9 +133,10 @@ public void testQos429_retryAfter_invalid() {
         assertThat(decoder.isError(response)).isTrue();
 
         RuntimeException result = decoder.decode(response);
-        assertThat(result)
-                .isInstanceOfSatisfying(QosException.Throttle.class, exception -> assertThat(exception.getRetryAfter())
-                        .isEmpty());
+        assertThat(result).isInstanceOfSatisfying(QosException.Throttle.class, exception -> {
+            assertThat(exception.getReason()).isEqualTo(ErrorDecoder.QOS_REASON);
+            assertThat(exception.getRetryAfter()).isEmpty();
+        });
     }
 
     @Test