From 18637d989d18cfbcc620d5030cca7cc569a2f2bf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=EA=B0=95=EC=A7=80=EC=9B=90?= <kangji149@naver.com>
Date: Mon, 28 Oct 2024 22:44:49 +0900
Subject: [PATCH] Fix gRPC Retry Mechanism for Unsuccessful HTTP Responses

---
 .../okhttp/internal/OkHttpGrpcSender.java     |  7 +-
 .../okhttp/internal/OkHttpGrpcSenderTest.java | 65 +++++++++++++++++++
 2 files changed, 69 insertions(+), 3 deletions(-)
 create mode 100644 exporters/sender/okhttp/src/test/java/io/opentelemetry/exporter/sender/okhttp/internal/OkHttpGrpcSenderTest.java

diff --git a/exporters/sender/okhttp/src/main/java/io/opentelemetry/exporter/sender/okhttp/internal/OkHttpGrpcSender.java b/exporters/sender/okhttp/src/main/java/io/opentelemetry/exporter/sender/okhttp/internal/OkHttpGrpcSender.java
index 8776762b62b..b873c4c2553 100644
--- a/exporters/sender/okhttp/src/main/java/io/opentelemetry/exporter/sender/okhttp/internal/OkHttpGrpcSender.java
+++ b/exporters/sender/okhttp/src/main/java/io/opentelemetry/exporter/sender/okhttp/internal/OkHttpGrpcSender.java
@@ -201,14 +201,15 @@ public CompletableResultCode shutdown() {
 
   /** Whether response is retriable or not. */
   public static boolean isRetryable(Response response) {
-    // Only retry on gRPC codes which will always come with an HTTP success
-    if (!response.isSuccessful()) {
+
+    String grpcStatus = grpcStatus(response);
+
+    if (grpcStatus == null) {
       return false;
     }
 
     // We don't check trailers for retry since retryable error codes always come with response
     // headers, not trailers, in practice.
-    String grpcStatus = response.header(GRPC_STATUS);
     return RetryUtil.retryableGrpcStatusCodes().contains(grpcStatus);
   }
 
diff --git a/exporters/sender/okhttp/src/test/java/io/opentelemetry/exporter/sender/okhttp/internal/OkHttpGrpcSenderTest.java b/exporters/sender/okhttp/src/test/java/io/opentelemetry/exporter/sender/okhttp/internal/OkHttpGrpcSenderTest.java
new file mode 100644
index 00000000000..8e12e3a7e7a
--- /dev/null
+++ b/exporters/sender/okhttp/src/test/java/io/opentelemetry/exporter/sender/okhttp/internal/OkHttpGrpcSenderTest.java
@@ -0,0 +1,65 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.exporter.sender.okhttp.internal;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import io.opentelemetry.exporter.internal.RetryUtil;
+import okhttp3.Protocol;
+import okhttp3.Request;
+import okhttp3.Response;
+import org.junit.jupiter.api.Test;
+
+public class OkHttpGrpcSenderTest {
+
+  @Test
+  public void isRetryable_SuccessfulResponse() {
+    Response response =
+        new Response.Builder()
+            .request(new Request.Builder().url("http://localhost/").build())
+            .protocol(Protocol.HTTP_2)
+            .code(200)
+            .message("OK")
+            .header("grpc-status", "0") // OK
+            .build();
+
+    boolean isRetryable = OkHttpGrpcSender.isRetryable(response);
+    assertFalse(isRetryable);
+  }
+
+  @Test
+  public void isRetryable_RetryableGrpcStatus() {
+    for (String grpcStatus : RetryUtil.retryableGrpcStatusCodes()) {
+      Response response =
+          new Response.Builder()
+              .request(new Request.Builder().url("http://localhost/").build())
+              .protocol(Protocol.HTTP_2)
+              .code(503)
+              .message("Service Unavailable")
+              .header("grpc-status", grpcStatus)
+              .build();
+
+      boolean isRetryable = OkHttpGrpcSender.isRetryable(response);
+      assertTrue(isRetryable);
+    }
+  }
+
+  @Test
+  public void isRetryable_NonRetryableGrpcStatus() {
+    Response response =
+        new Response.Builder()
+            .request(new Request.Builder().url("http://localhost/").build())
+            .protocol(Protocol.HTTP_2)
+            .code(400)
+            .message("Bad Request")
+            .header("grpc-status", "3") // INVALID_ARGUMENT
+            .build();
+
+    boolean isRetryable = OkHttpGrpcSender.isRetryable(response);
+    assertFalse(isRetryable);
+  }
+}