fix(gpu): contain errors in unwatched compositor work (#4894)

* fix(gpu): contain errors in unwatched compositor work

* test(gpu): gate recovery fixtures on format capabilities

---------

Co-authored-by: haranishi <[email protected]>

crates/gpu/src/lib.rs

@@ -552,17 +552,18 @@ fn field_reaches(progs: &[fx::MapProgram]) -> HashMap<FieldKind, i32> {
     want
 }
 
-/// Runs `f` inside validation and internal-error scopes and returns the first error raised.
+/// Runs `f` inside GPU error scopes and returns the first error raised.
 /// Native wgpu reports errors synchronously; where they only arrive asynchronously (WebGPU) none
 /// is reported.
 fn first_error<T>(device: &wgpu::Device, f: impl FnOnce() -> T) -> (T, Option<wgpu::Error>) {
     let validation = device.push_error_scope(wgpu::ErrorFilter::Validation);
     let internal = device.push_error_scope(wgpu::ErrorFilter::Internal);
+    let memory = device.push_error_scope(wgpu::ErrorFilter::OutOfMemory);
     let out = f();
     let mut cx = std::task::Context::from_waker(std::task::Waker::noop());
     let mut error = None;
     // Scopes pop in reverse order.
-    for pop in [internal.pop(), validation.pop()] {
+    for pop in [memory.pop(), internal.pop(), validation.pop()] {
         if let std::task::Poll::Ready(Some(e)) = std::pin::pin!(pop).poll(&mut cx) {
             error.get_or_insert(e);
         }
@@ -570,6 +571,21 @@ fn first_error<T>(device: &wgpu::Device, f: impl FnOnce() -> T) -> (T, Option<wg
     (out, error)
 }
 
+/// Unwatched native callers need GPU errors to become `Unsupported` instead of reaching wgpu's
+/// default panic handler. Watched callers keep their fault callbacks, which stop work promptly
+/// and update the health flag shared with the canvas.
+fn scoped_work<T>(device: &wgpu::Device, watched: bool, f: impl FnOnce() -> Result<T, Unsupported>) -> Result<T, Unsupported> {
+    // WebGPU errors arrive asynchronously; its callers still use the health callbacks.
+    if watched || cfg!(target_arch = "wasm32") {
+        return f();
+    }
+    let (result, error) = first_error(device, f);
+    match error {
+        Some(error) => Err(Unsupported(format!("GPU work failed: {error}"))),
+        None => result,
+    }
+}
+
 impl Compositor {
     /// Create a compositor with the default accumulation format ([`ACC_FORMAT`]).
     pub fn new(device: &wgpu::Device) -> Self {
@@ -902,16 +918,16 @@ impl Compositor {
         region: Rect,
         mut sink: impl FnMut(&mut wgpu::CommandEncoder, ChunkOut<'_>),
     ) -> Result<Stats, Unsupported> {
-        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("pc_compose") });
-        let stats = match self.encode_inner(device, queue, &mut encoder, doc, region, &mut sink, true) {
-            Ok(stats) => stats,
-            Err(error) => {
-                self.prefix.clear();
-                return Err(error);
-            }
-        };
-        queue.submit([encoder.finish()]);
-        Ok(stats)
+        let result = scoped_work(device, self.health.is_some(), || {
+            let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("pc_compose") });
+            let stats = self.encode_inner(device, queue, &mut encoder, doc, region, &mut sink, true)?;
+            queue.submit([encoder.finish()]);
+            Ok(stats)
+        });
+        if result.is_err() {
+            self.discard_failed_work();
+        }
+        result
     }
 
     /// Like [`Self::render`] but records into `encoder` without submitting.
@@ -924,7 +940,22 @@ impl Compositor {
         region: Rect,
         sink: &mut dyn FnMut(&mut wgpu::CommandEncoder, ChunkOut<'_>),
     ) -> Result<Stats, Unsupported> {
-        self.encode_inner(device, queue, encoder, doc, region, sink, false)
+        let result = scoped_work(device, self.health.is_some(), || self.encode_inner(device, queue, encoder, doc, region, sink, false));
+        if result.is_err() {
+            self.discard_failed_work();
+        }
+        result
+    }
+
+    fn discard_failed_work(&mut self) {
+        // Cache metadata advances while recording. A failed encoder may never execute its
+        // texture uploads or effect maps, so none of those cached pixels can be trusted.
+        self.prefix.clear();
+        self.residents.clear();
+        self.fx.clear();
+        self.patterns.clear();
+        self.temps.clear();
+        self.pool.clear();
     }
 
     #[allow(clippy::too_many_arguments)]

crates/gpu/tests/device_loss.rs

@@ -81,3 +81,181 @@ fn health_wait_on_a_destroyed_device_is_safe() {
     assert!(!health.is_ok());
     assert!(!health.wait(&device, None));
 }
+
+fn assert_unwatched_destroyed_device(encode_only: bool) {
+    let _lock = GPU_LOCK.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
+    let Some((adapter, device, queue)) = device() else { return };
+    let mut comp = Compositor::try_new_with_format(&device, Compositor::preferred_acc_format(&adapter)).unwrap();
+    let d = Document::with_background("synthetic", Size::new(1, 1), ColorMode::Rgb, SampleType::U8, Color::rgba(0.9, 0.9, 0.9, 1.0));
+    assert!(render_to_vec(&mut comp, &device, &queue, &d, d.bounds()).is_ok());
+    let mut encoder = device.create_command_encoder(&Default::default());
+    device.destroy();
+    let mut changed = d.clone();
+    changed.layers[0].opacity = 0.4;
+    let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+        if encode_only {
+            comp.encode(&device, &queue, &mut encoder, &changed, changed.bounds(), &mut |_, _| {})
+        } else {
+            comp.render(&device, &queue, &changed, changed.bounds(), |_, _| {})
+        }
+    }));
+    assert!(result.is_ok(), "an unwatched device failure must return Unsupported instead of panicking");
+    let error = result.unwrap().expect_err("a destroyed device must not report successful GPU work");
+    assert!(error.0.contains("GPU work failed"), "{error}");
+}
+
+#[test]
+fn an_unwatched_destroyed_device_returns_an_error_without_panicking() {
+    assert_unwatched_destroyed_device(false);
+}
+
+#[test]
+fn an_unwatched_encode_on_a_destroyed_device_returns_an_error_without_panicking() {
+    assert_unwatched_destroyed_device(true);
+}
+
+#[test]
+fn loss_during_an_unwatched_render_returns_an_error_without_panicking() {
+    let _lock = GPU_LOCK.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
+    let Some((adapter, device, queue)) = device() else { return };
+    let mut comp = Compositor::try_new_with_format(&device, Compositor::preferred_acc_format(&adapter)).unwrap();
+    comp.set_texture_limit(256);
+    let d = Document::with_background("synthetic", Size::new(512, 1), ColorMode::Rgb, SampleType::U8, Color::rgba(0.9, 0.9, 0.9, 1.0));
+    let mut chunks = 0;
+    let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+        comp.render(&device, &queue, &d, d.bounds(), |_, _| {
+            chunks += 1;
+            if chunks == 1 {
+                // A driver reset can arrive after the entry-point health check too.
+                device.destroy();
+            }
+        })
+    }));
+    assert!(chunks > 0, "the injection must occur after GPU work began");
+    assert!(result.is_ok(), "device loss mid-render must return Unsupported instead of panicking");
+    assert!(result.unwrap().is_err(), "a partially failed render must not report success");
+}
+
+#[test]
+fn a_captured_render_error_preserves_the_device_for_a_later_render() {
+    let _lock = GPU_LOCK.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
+    let Some((adapter, device, queue)) = device() else { return };
+    let mut comp = Compositor::try_new_with_format(&device, Compositor::preferred_acc_format(&adapter)).unwrap();
+    let d = Document::with_background("synthetic", Size::new(1, 1), ColorMode::Rgb, SampleType::U8, Color::rgba(0.9, 0.9, 0.9, 1.0));
+    let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+        comp.render(&device, &queue, &d, d.bounds(), |_, _| {
+            // Reject a buffer above the device's valid capability without actually allocating it.
+            let _ = device.create_buffer(&wgpu::BufferDescriptor {
+                label: Some("synthetic_invalid_buffer"),
+                size: device.limits().max_buffer_size + 1,
+                usage: wgpu::BufferUsages::COPY_DST,
+                mapped_at_creation: false,
+            });
+        })
+    }));
+    assert!(result.is_ok(), "GPU validation errors must return Unsupported instead of panicking");
+    let error = result.unwrap().expect_err("an invalid allocation must fail this render");
+    assert!(error.0.contains("synthetic_invalid_buffer"), "{error}");
+    let out = render_to_vec(&mut comp, &device, &queue, &d, d.bounds()).expect("the next fitting render keeps working");
+    let expected = photocraft_compose::flatten(&d);
+    for (cpu, gpu) in expected.px[0].iter().zip(out[0]) {
+        assert!((cpu - gpu).abs() <= 1.0 / 255.0, "the healthy GPU control differs: {cpu} vs {gpu}");
+    }
+}
+
+fn assert_recovery_after_invalid_encoder(encode_only: bool) {
+    let _lock = GPU_LOCK.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
+    let Some((adapter, device, queue)) = device() else { return };
+    let mut comp = Compositor::try_new_with_format(&device, Compositor::preferred_acc_format(&adapter)).unwrap();
+    let mut d = doc();
+    d.layers[1].effects.items = vec![photocraft_doc::Effect::default_drop_shadow()];
+    // Pipeline creation errors are not proof of missing hardware capabilities: check the
+    // adapter and the features actually enabled on this device before allowing a skip.
+    let required = wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::TEXTURE_BINDING;
+    let features = device.features();
+    for format in [wgpu::TextureFormat::R32Float, wgpu::TextureFormat::R16Float] {
+        let adapter_usages = adapter.get_texture_format_features(format).allowed_usages;
+        let guaranteed_usages = format.guaranteed_format_features(features).allowed_usages;
+        if !adapter_usages.contains(required)
+            || (!features.contains(wgpu::Features::TEXTURE_ADAPTER_SPECIFIC_FORMAT_FEATURES) && !guaranteed_usages.contains(required))
+        {
+            eprintln!(
+                "skipping effect-map recovery test: {:?} {format:?} lacks {required:?}; adapter usages={adapter_usages:?}, guaranteed usages={guaranteed_usages:?}, device features={features:?}",
+                adapter.get_info().backend,
+            );
+            return;
+        }
+    }
+    comp.supports(&d).expect("a capable adapter must support the recovery fixture's effect-map pipelines");
+    let src = device.create_buffer(&wgpu::BufferDescriptor { label: None, size: 4, usage: wgpu::BufferUsages::COPY_SRC, mapped_at_creation: false });
+    let dst = device.create_buffer(&wgpu::BufferDescriptor { label: None, size: 4, usage: wgpu::BufferUsages::COPY_DST, mapped_at_creation: false });
+    let mut injected = false;
+    let mut sink = |encoder: &mut wgpu::CommandEncoder, _: photocraft_gpu::ChunkOut<'_>| {
+        injected = true;
+        // Misaligned copies invalidate the encoder, discarding all recorded effect-map work.
+        encoder.copy_buffer_to_buffer(&src, 0, &dst, 0, 1);
+        if encode_only {
+            // encode does not finish its caller-owned encoder; force an immediate error too.
+            let _ = device.create_buffer(&wgpu::BufferDescriptor {
+                label: Some("synthetic_invalid_buffer"),
+                size: device.limits().max_buffer_size + 1,
+                usage: wgpu::BufferUsages::COPY_DST,
+                mapped_at_creation: false,
+            });
+        }
+    };
+    let result = if encode_only {
+        let mut encoder = device.create_command_encoder(&Default::default());
+        comp.encode(&device, &queue, &mut encoder, &d, d.bounds(), &mut sink)
+    } else {
+        comp.render(&device, &queue, &d, d.bounds(), sink)
+    };
+    assert!(injected, "the invalid encoder injection must reach the sink");
+    let error = result.expect_err("invalid encoder must fail");
+    assert!(error.0.contains("GPU work failed"), "the injected GPU error must be captured: {error}");
+    let recovered = render_to_vec(&mut comp, &device, &queue, &d, d.bounds()).unwrap();
+    let mut fresh = Compositor::try_new_with_format(&device, Compositor::preferred_acc_format(&adapter)).unwrap();
+    let expected = render_to_vec(&mut fresh, &device, &queue, &d, d.bounds()).unwrap();
+    for (a, b) in recovered.iter().zip(expected.iter()) {
+        for (a, b) in a.iter().zip(b.iter()) {
+            assert!((a - b).abs() <= 1.0 / 255.0, "recovered cached pixels differ: {a} vs {b}");
+        }
+    }
+}
+
+#[test]
+fn render_recovers_effect_maps_after_invalid_encoder() {
+    assert_recovery_after_invalid_encoder(false);
+}
+
+#[test]
+fn encode_recovers_effect_maps_after_invalid_encoder() {
+    assert_recovery_after_invalid_encoder(true);
+}
+
+#[test]
+fn loss_during_intermediate_unwatched_submissions_returns_an_error() {
+    let _lock = GPU_LOCK.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
+    let Some((adapter, device, queue)) = device() else { return };
+    let mut comp = Compositor::try_new_with_format(&device, Compositor::preferred_acc_format(&adapter)).unwrap();
+    comp.set_texture_limit(256);
+    let mut d = Document::with_background("synthetic", Size::new(2048, 1), ColorMode::Rgb, SampleType::U8, Color::rgba(0.9, 0.9, 0.9, 1.0));
+    for _ in 0..300 {
+        let mut layer = d.layers[0].clone();
+        layer.id = photocraft_doc::LayerId::fresh();
+        layer.opacity = 0.5;
+        d.layers.push(layer);
+    }
+    let mut chunks = 0;
+    let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+        comp.render(&device, &queue, &d, d.bounds(), |_, _| {
+            chunks += 1;
+            if chunks == 1 {
+                device.destroy();
+            }
+        })
+    }));
+    assert!(chunks > 0);
+    assert!(result.is_ok(), "intermediate submit polling must not escape as a panic");
+    assert!(result.unwrap().is_err());
+}