From 26da0c096d4a48a0d7f1d292b7382d7b2b24715e Mon Sep 17 00:00:00 2001
From: kronael <ondrej@vostal.net>
Date: Mon, 3 Jun 2024 09:26:24 +0200
Subject: [PATCH] Add panic hook to exit on panic in client (#44)

* Add panic hook to exit on panic to client

* Run cargo fmt
---
 Makefile                 |  9 ++++++++-
 client/main.rs           | 13 +++++++++++++
 server/balancer.rs       |  5 +++--
 server/solana_service.rs |  2 +-
 4 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/Makefile b/Makefile
index be3d364..e6ebae8 100644
--- a/Makefile
+++ b/Makefile
@@ -7,7 +7,7 @@ help:
 	@echo "    image           builds $(name) docker image"
 	@echo "    help            show this help"
 
-.PHONY: build-all build-server build-server-release run-server clean run-client run-client-local
+.PHONY: build-all build-server build-server-release run-server clean run-client run-client-local run-client-rpc-devnet run-client-rpc-mainnet run-client-blackhole
 .DEFAULT_GOAL := build-all
 
 CERT_DIR = certs
@@ -59,6 +59,13 @@ run-server-local: build-server
 		--tls-grpc-server-key     ./certs/localhost.key \
 		--tls-grpc-ca-cert        ./certs/ca.cert \
 
+run-client: build-client
+	cargo run --bin mtx-client -- \
+		--tls-grpc-ca-cert     ./certs/ca.cert \
+		--grpc-urls-file       ./client.yml \
+		--tls-grpc-client-key  ./certs/client.$(client).key \
+		--tls-grpc-client-cert ./certs/client.$(client).cert \
+
 run-client-rpc-devnet: build-client
 	cargo run --bin mtx-client -- \
 		--tls-grpc-ca-cert     ./certs/ca.cert \
diff --git a/client/main.rs b/client/main.rs
index 5f02573..7343679 100644
--- a/client/main.rs
+++ b/client/main.rs
@@ -11,6 +11,8 @@ use forwarder::ForwardedTransaction;
 use log::{error, info};
 use signal_hook_tokio::Signals;
 use solana_sdk::signature::read_keypair_file;
+use std::panic;
+use std::process;
 use structopt::StructOpt;
 use tokio::{
     sync::{mpsc::UnboundedSender, RwLock},
@@ -64,8 +66,19 @@ struct Params {
     throttle_parallel: usize,
 }
 
+fn setup_panic_hook() {
+    let hook = panic::take_hook();
+    panic::set_hook(Box::new(move |info| {
+        hook(info);
+        error!("panic_hook forcing exit");
+        process::exit(1);
+    }));
+}
+
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    setup_panic_hook();
+
     env_logger::Builder::from_env(Env::default().default_filter_or("info")).init();
 
     let params = Params::from_args();
diff --git a/server/balancer.rs b/server/balancer.rs
index 727707a..bcb692c 100644
--- a/server/balancer.rs
+++ b/server/balancer.rs
@@ -1,7 +1,7 @@
 use crate::grpc_server::{self, build_tx_message_envelope};
 use crate::metrics;
 use crate::solana_service::{get_tpu_by_identity, leaders_stream};
-use crate::{N_CONSUMERS, N_COPIES, NODES_REFRESH_SECONDS};
+use crate::{NODES_REFRESH_SECONDS, N_CONSUMERS, N_COPIES};
 use jsonrpc_http_server::*;
 use log::{error, info};
 use rand::rngs::StdRng;
@@ -207,7 +207,8 @@ pub fn balancer_updater(
     );
 
     let mut refresh_cluster_nodes_hint = Box::pin(
-        tokio_stream::iter(std::iter::repeat(())).throttle(tokio::time::Duration::from_secs(NODES_REFRESH_SECONDS)),
+        tokio_stream::iter(std::iter::repeat(()))
+            .throttle(tokio::time::Duration::from_secs(NODES_REFRESH_SECONDS)),
     );
 
     let mut tpu_by_identity = Default::default();
diff --git a/server/solana_service.rs b/server/solana_service.rs
index 03d6b44..b2114db 100644
--- a/server/solana_service.rs
+++ b/server/solana_service.rs
@@ -1,5 +1,5 @@
-use crate::{ N_LEADERS, LEADER_REFRESH_SECONDS };
 use crate::{metrics, rpc_server::Mode};
+use crate::{LEADER_REFRESH_SECONDS, N_LEADERS};
 use log::{debug, error, info};
 use solana_client::{
     nonblocking::pubsub_client::PubsubClient, rpc_client::RpcClient,