From 513ec31d1eb0dc0f1dd6c78a6e2de584bcef0e21 Mon Sep 17 00:00:00 2001
From: Trent Nelson <trent@solana.com>
Date: Wed, 17 Feb 2021 01:51:19 -0700
Subject: [PATCH] SoM: Retry HTTP 502/504

---
 Cargo.lock                | 1 +
 stake-o-matic/Cargo.toml  | 1 +
 stake-o-matic/src/main.rs | 8 +++++++-
 3 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/Cargo.lock b/Cargo.lock
index b215a343b4..a2c049a8ab 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -5234,6 +5234,7 @@ version = "1.6.0"
 dependencies = [
  "clap",
  "log 0.4.11",
+ "reqwest 0.10.8",
  "semver 0.9.0",
  "serde_yaml",
  "solana-clap-utils",
diff --git a/stake-o-matic/Cargo.toml b/stake-o-matic/Cargo.toml
index a7c1ed441a..f9da2257f6 100644
--- a/stake-o-matic/Cargo.toml
+++ b/stake-o-matic/Cargo.toml
@@ -11,6 +11,7 @@ version = "1.6.0"
 [dependencies]
 clap = "2.33.0"
 log = "0.4.11"
+reqwest = { version = "0.10.8", default-features = false, features = ["blocking", "rustls-tls", "json"] }
 semver = "0.9.0"
 serde_yaml = "0.8.13"
 solana-clap-utils = { path = "../clap-utils", version = "1.6.0" }
diff --git a/stake-o-matic/src/main.rs b/stake-o-matic/src/main.rs
index 58c53479f3..2accf82e81 100644
--- a/stake-o-matic/src/main.rs
+++ b/stake-o-matic/src/main.rs
@@ -5,6 +5,7 @@ use {
         ArgMatches,
     },
     log::*,
+    reqwest::StatusCode,
     solana_clap_utils::{
         input_parsers::{keypair_of, pubkey_of},
         input_validators::{
@@ -371,7 +372,12 @@ where
             ..
         }) = result
         {
-            if reqwest_error.is_timeout() && retries > 0 {
+            let can_retry = reqwest_error.is_timeout()
+                || reqwest_error
+                    .status()
+                    .map(|s| s == StatusCode::BAD_GATEWAY || s == StatusCode::GATEWAY_TIMEOUT)
+                    .unwrap_or(false);
+            if can_retry && retries > 0 {
                 info!("RPC request timeout, {} retries remaining", retries);
                 retries -= 1;
                 continue;