3 files changed, 33 insertions, 2 deletions
diff --git a/pingora-proxy/src/lib.rs b/pingora-proxy/src/lib.rs
index f13a1f7..731cdb4 100644
--- a/pingora-proxy/src/lib.rs
+++ b/pingora-proxy/src/lib.rs
@@ -535,6 +535,10 @@ impl<SV> HttpProxy<SV> {
                 if !proxy_to_upstream {
                     // The hook can choose to write its own response, but if it doesn't, we respond
                     // with a generic 502
+                    if session.cache.enabled() {
+                        // drop the cache lock that this request may be holding onto
+                        session.cache.disable(NoCacheReason::DeclinedToUpstream);
+                    }
                     if session.response_written().is_none() {
                         match session.write_response_header_ref(&BAD_GATEWAY).await {
                             Ok(()) => {}
@@ -557,6 +561,10 @@ impl<SV> HttpProxy<SV> {
                 /* else continue */
             }
             Err(e) => {
+                if session.cache.enabled() {
+                    session.cache.disable(NoCacheReason::InternalError);
+                }
+
                 self.handle_error(
                     &mut session,
                     &mut ctx,
@@ -621,7 +629,14 @@ impl<SV> HttpProxy<SV> {
 
         if let Some(e) = final_error.as_ref() {
             // If we have errored and are still holding a cache lock, release it.
-            session.cache.disable(NoCacheReason::InternalError);
+            if session.cache.enabled() {
+                let reason = if *e.esource() == ErrorSource::Upstream {
+                    NoCacheReason::UpstreamError
+                } else {
+                    NoCacheReason::InternalError
+                };
+                session.cache.disable(reason);
+            }
             let status = self.inner.fail_to_proxy(&mut session, e, &mut ctx).await;
 
             // final error will have > 0 status unless downstream connection is dead
diff --git a/pingora-proxy/src/proxy_cache.rs b/pingora-proxy/src/proxy_cache.rs
index 4369ec8..8957bc7 100644
--- a/pingora-proxy/src/proxy_cache.rs
+++ b/pingora-proxy/src/proxy_cache.rs
@@ -667,8 +667,18 @@ impl<SV> HttpProxy<SV> {
                         None,
                         None,
                     );
-                    self.inner
+                    if self
+                        .inner
                         .should_serve_stale(session, ctx, Some(&http_status_error))
+                    {
+                        // no more need to keep the write lock
+                        session
+                            .cache
+                            .release_write_lock(NoCacheReason::UpstreamError);
+                        true
+                    } else {
+                        false
+                    }
                 } else {
                     false // not 304, not stale if error status code
                 }
@@ -712,6 +722,11 @@ impl<SV> HttpProxy<SV> {
             self.inner.request_summary(session, ctx)
         );
 
+        // no more need to hang onto the cache lock
+        session
+            .cache
+            .release_write_lock(NoCacheReason::UpstreamError);
+
         Some(self.proxy_cache_hit(session, ctx).await)
     }
 
diff --git a/pingora-proxy/src/proxy_h2.rs b/pingora-proxy/src/proxy_h2.rs
index ba832b2..a27c559 100644
--- a/pingora-proxy/src/proxy_h2.rs
+++ b/pingora-proxy/src/proxy_h2.rs
@@ -409,6 +409,7 @@ impl<SV> HttpProxy<SV> {
                     .cache_http_task(session, &task, ctx, serve_from_cache)
                     .await
                 {
+                    session.cache.disable(NoCacheReason::StorageError);
                     if serve_from_cache.is_miss_body() {
                         // if the response stream cache body during miss but write fails, it has to
                         // give up the entire request