From 13ed1e32b4f978151ac49236756679b99d383a7e Mon Sep 17 00:00:00 2001
From: Bojan Serafimov <bojan@pop-os.localdomain>
Date: Sun, 22 May 2022 16:51:12 -0400
Subject: [PATCH] Handle panic

---
 pageserver/src/jobs/scheduler.rs | 18 +++++++++++++-----
 pageserver/src/jobs/worker.rs    | 13 ++++++++++---
 2 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/pageserver/src/jobs/scheduler.rs b/pageserver/src/jobs/scheduler.rs
index 3ac9a81576..67e6e0b4d3 100644
--- a/pageserver/src/jobs/scheduler.rs
+++ b/pageserver/src/jobs/scheduler.rs
@@ -55,11 +55,19 @@ impl<J: Job> Sched<J> {
                     report = self.report.1.recv() => {
                         // Reschedule job to run again
                         let send_work = self.work.0.clone();
-                        let job = report.unwrap().for_job;
-                        tokio::spawn(async move {
-                            sleep(Duration::from_millis(10)).await;
-                            send_work.send(job).await.unwrap();
-                        });
+                        let report = report.unwrap();
+                        let job = report.for_job;
+                        match report.result {
+                            Ok(()) => {
+                                tokio::spawn(async move {
+                                    sleep(Duration::from_millis(10)).await;
+                                    send_work.send(job).await.unwrap();
+                                });
+                            },
+                            Err(e) => {
+                                println!("task panicked");
+                            }
+                        }
                     }
                 }
             }
diff --git a/pageserver/src/jobs/worker.rs b/pageserver/src/jobs/worker.rs
index a535fa596c..e0f56fa931 100644
--- a/pageserver/src/jobs/worker.rs
+++ b/pageserver/src/jobs/worker.rs
@@ -3,6 +3,9 @@
 //!
 use crate::thread_mgr::shutdown_watcher;
 use tokio::sync::mpsc::{Sender, channel};
+use std::any::Any;
+use std::panic::AssertUnwindSafe;
+use std::panic::catch_unwind;
 
 pub trait Job: std::fmt::Debug + Send + 'static {
     fn run(&self);
@@ -14,6 +17,7 @@ pub struct Worker<J: Job>(pub Sender<J>);
 #[derive(Debug)]
 pub struct Report<J: Job> {
     pub for_job: J,
+    pub result: Result<(), Box<dyn Any + Send>>
 }
 
 pub fn run_worker<J: Job>(enlist: Sender<Worker<J>>, report: Sender<Report<J>>) -> anyhow::Result<()> {
@@ -31,13 +35,16 @@ pub fn run_worker<J: Job>(enlist: Sender<Worker<J>>, report: Sender<Report<J>>)
                 _ = shutdown_watcher => break,
                 j = get_work.recv() => {
                     if let Some(job) = j {
-                        job.run();
+                        let result = catch_unwind(AssertUnwindSafe(|| {
+                            job.run();
+                        }));
                         report.send(Report {
                             for_job: job,
+                            result: result,
                         }).await.unwrap();
                     } else {
-                        println!("fffffffff")
-                        // TODO ??
+                        // channel closed
+                        return;
                     }
                 }
             };