Support NeedData responses in SourceMap/SourceMapTrace ContentSources (…

…#3907)

This implements a new `ContentSourceProcessor` trait and
`WrappedContentSource` struct that allows our `ContentSource`s to
register some processor to transform the eventual `ContentSourceContent`
that an inner `ContentSource` returns. Yah, it's a bit of a mouthful
with lots of very similar sounding names.

Essentially, the old `SourceMapContentSource` and
`NextSourceMapTraceContentSource` wrapped some inner `ContentSource`,
and would request content from them, and process that content into a
source map JSON or trace JSON. But, they didn't implement the `NeedData`
response, so it only supported very primitive inner content sources.

This PR extracts that knowledge into a single `WrappedContentSource`,
which will recursively wrap every `ContentSourceResult` until we land on
a fully resolved `ContentSourceResult::Result`. At that point, it hands
off the work to a `WrappedGetContentSource`, which performs the
processing on whatever content is returned by the inner
`GetContentSourceContent`. If you can't tell yet, I'm making the
description intentionally verbose to highlight just how similar our
struct/trait names are.

Fixes WEB-614

crates/next-core/src/next_image/mod.rs

@@ -5,8 +5,8 @@ use turbo_tasks::{primitives::StringVc, Value};
 use turbopack_core::introspect::{Introspectable, IntrospectableVc};
 use turbopack_dev_server::source::{
     query::QueryValue, ContentSource, ContentSourceContent, ContentSourceData,
-    ContentSourceDataFilter, ContentSourceDataVary, ContentSourceResult, ContentSourceResultVc,
-    ContentSourceVc, NeededData, ProxyResult,
+    ContentSourceDataFilter, ContentSourceDataVary, ContentSourceResultVc, ContentSourceVc,
+    NeededData, ProxyResult,
 };
 
 /// Serves, resizes, optimizes, and re-encodes images to be used with
@@ -65,11 +65,7 @@ impl ContentSource for NextImageContentSource {
         // TODO: consume the assets, resize and reduce quality, re-encode into next-gen
         // formats.
         if let Some(path) = url.strip_prefix('/') {
-            let asset = this.asset_source.get(path, Default::default());
-            let inner = asset.await?;
-            if let ContentSourceResult::Result { .. } = &*inner {
-                return Ok(asset);
-            }
+            return Ok(this.asset_source.get(path, Default::default()));
         }
 
         // TODO: This should be downloaded by the server, and resized, etc.

crates/turbopack-dev-server/src/source/mod.rs

@@ -10,6 +10,7 @@ pub mod router;
 pub mod source_maps;
 pub mod specificity;
 pub mod static_assets;
+pub mod wrapping_source;
 
 use std::{collections::BTreeSet, sync::Arc};
 

crates/turbopack-dev-server/src/source/source_maps.rs

@@ -1,4 +1,5 @@
 use anyhow::Result;
+use mime::APPLICATION_JSON;
 use turbo_tasks::{primitives::StringVc, Value};
 use turbo_tasks_fs::File;
 use turbopack_core::{
@@ -8,11 +9,12 @@ use turbopack_core::{
 };
 
 use super::{
-    query::QueryValue, ContentSource, ContentSourceContent, ContentSourceContentVc,
-    ContentSourceData, ContentSourceDataFilter, ContentSourceDataVary, ContentSourceResult,
-    ContentSourceResultVc, ContentSourceVc, NeededData,
+    query::QueryValue,
+    wrapping_source::{ContentSourceProcessor, ContentSourceProcessorVc, WrappedContentSourceVc},
+    ContentSource, ContentSourceContent, ContentSourceContentVc, ContentSourceData,
+    ContentSourceDataFilter, ContentSourceDataVary, ContentSourceResultVc, ContentSourceVc,
+    NeededData,
 };
-use crate::source::GetContentSourceContent;
 
 /// SourceMapContentSource allows us to serve full source maps, and individual
 /// sections of source maps, of any found asset in the graph without adding
@@ -66,55 +68,77 @@ impl ContentSource for SourceMapContentSource {
         };
 
         let id = match query.get("id") {
-            Some(QueryValue::String(s)) => Some(s),
+            Some(QueryValue::String(s)) => Some(s.clone()),
             _ => None,
         };
 
-        let this = self_vc.await?;
-        let result = this.asset_source.get(pathname, Default::default()).await?;
-        let content = match &*result {
-            ContentSourceResult::Result { get_content, .. } => {
-                get_content.get(Default::default()).await?
-            }
-            _ => return Ok(ContentSourceResultVc::not_found()),
-        };
-        let file = match &*content {
+        let wrapped = WrappedContentSourceVc::new(
+            self_vc.await?.asset_source,
+            SourceMapContentProcessorVc::new(id).into(),
+        );
+        Ok(wrapped
+            .as_content_source()
+            .get(pathname, Default::default()))
+    }
+}
+
+#[turbo_tasks::value_impl]
+impl Introspectable for SourceMapContentSource {
+    #[turbo_tasks::function]
+    fn ty(&self) -> StringVc {
+        StringVc::cell("source map content source".to_string())
+    }
+
+    #[turbo_tasks::function]
+    fn details(&self) -> StringVc {
+        StringVc::cell("serves chunk and chunk item source maps".to_string())
+    }
+}
+
+/// Processes the eventual [ContentSourceContent] to transform it into a source
+/// map JSON content.
+#[turbo_tasks::value]
+pub struct SourceMapContentProcessor {
+    /// An optional section id to use when generating the map. Specifying a
+    /// section id will only output that section. Otherwise, it prints the
+    /// full source map.
+    id: Option<String>,
+}
+
+#[turbo_tasks::value_impl]
+impl SourceMapContentProcessorVc {
+    #[turbo_tasks::function]
+    fn new(id: Option<String>) -> Self {
+        SourceMapContentProcessor { id }.cell()
+    }
+}
+
+#[turbo_tasks::value_impl]
+impl ContentSourceProcessor for SourceMapContentProcessor {
+    #[turbo_tasks::function]
+    async fn process(&self, content: ContentSourceContentVc) -> Result<ContentSourceContentVc> {
+        let file = match &*content.await? {
             ContentSourceContent::Static(static_content) => static_content.await?.content,
-            _ => return Ok(ContentSourceResultVc::not_found()),
+            _ => return Ok(ContentSourceContentVc::not_found()),
         };
 
         let gen = match GenerateSourceMapVc::resolve_from(file).await? {
             Some(f) => f,
-            None => return Ok(ContentSourceResultVc::not_found()),
+            None => return Ok(ContentSourceContentVc::not_found()),
         };
 
-        let sm = if let Some(id) = id {
+        let sm = if let Some(id) = &self.id {
             let section = gen.by_section(id).await?;
             match &*section {
                 Some(sm) => *sm,
-                None => return Ok(ContentSourceResultVc::not_found()),
+                None => return Ok(ContentSourceContentVc::not_found()),
             }
         } else {
             gen.generate_source_map()
         };
-        let content = sm.to_rope().await?;
-
-        let asset = AssetContentVc::from(File::from(content));
-        Ok(ContentSourceResultVc::exact(
-            ContentSourceContentVc::static_content(asset.into()).into(),
-        ))
-    }
-}
 
-#[turbo_tasks::value_impl]
-impl Introspectable for SourceMapContentSource {
-    #[turbo_tasks::function]
-    fn ty(&self) -> StringVc {
-        StringVc::cell("source map content source".to_string())
-    }
-
-    #[turbo_tasks::function]
-    fn details(&self) -> StringVc {
-        StringVc::cell("serves chunk and chunk item source maps".to_string())
+        let content = sm.to_rope().await?;
+        let asset = AssetContentVc::from(File::from(content).with_content_type(APPLICATION_JSON));
+        Ok(ContentSourceContentVc::static_content(asset.into()))
     }
 }

crates/turbopack-dev-server/src/source/wrapping_source.rs

@@ -0,0 +1,122 @@
+use anyhow::Result;
+use turbo_tasks::Value;
+
+use super::{
+    ContentSource, ContentSourceContentVc, ContentSourceData, ContentSourceDataVaryVc,
+    ContentSourceResult, ContentSourceResultVc, ContentSourceVc, GetContentSourceContent,
+    GetContentSourceContentVc, NeededData,
+};
+
+/// A ContentSourceProcessor handles the final processing of an eventual
+/// [ContentSourceContent].
+///
+/// Used in conjunction with [WrappedContentSource], this allows a
+/// [ContentSource] implementation to easily register a final process step over
+/// some inner ContentSource's fully resolved [ContentSourceResult] and
+/// [ContentSourceContent].
+#[turbo_tasks::value_trait]
+pub trait ContentSourceProcessor {
+    fn process(&self, content: ContentSourceContentVc) -> ContentSourceContentVc;
+}
+
+/// A ContentSourceProcessor allows a [ContentSource] implementation to easily
+/// register a final process step over some inner ContentSource's fully resolved
+/// [ContentSourceResult] and [ContentSourceContent] without having to manually
+/// implement the NeedData resolution algorithm.
+///
+/// This is the first of 2 steps, implementing the wrapping of
+/// ContentSourceResult so that we can wrap the fully resolved result with our
+/// [WrappedGetContentSourceContent].
+#[turbo_tasks::value]
+pub struct WrappedContentSource {
+    inner: ContentSourceVc,
+    processor: ContentSourceProcessorVc,
+}
+
+#[turbo_tasks::value_impl]
+impl WrappedContentSourceVc {
+    #[turbo_tasks::function]
+    pub async fn new(inner: ContentSourceVc, processor: ContentSourceProcessorVc) -> Self {
+        WrappedContentSource { inner, processor }.cell()
+    }
+}
+
+#[turbo_tasks::value_impl]
+impl ContentSource for WrappedContentSource {
+    #[turbo_tasks::function]
+    async fn get(
+        &self,
+        path: &str,
+        data: Value<ContentSourceData>,
+    ) -> Result<ContentSourceResultVc> {
+        let res = self.inner.get(path, data);
+
+        Ok(match &*res.await? {
+            ContentSourceResult::NotFound => res,
+            ContentSourceResult::NeedData(needed) => {
+                // If the inner source needs more data, then we need to wrap the resuming source
+                // in a new wrapped processor. That way, whatever ContentSourceResult is
+                // returned when we resume can itself be wrapped, or be wrapped with a
+                // WrappedGetContentSourceContent.
+                ContentSourceResultVc::need_data(Value::new(NeededData {
+                    source: WrappedContentSourceVc::new(needed.source, self.processor).into(),
+                    path: needed.path.clone(),
+                    vary: needed.vary.clone(),
+                }))
+            }
+            ContentSourceResult::Result {
+                get_content,
+                specificity,
+            } => {
+                // If we landed on a result, then the resolution algorithm is complete. All
+                // that's left is to wrap the result's GetContentSourceContent
+                // with our own, so that we can process whatever content it
+                // returns.
+                ContentSourceResult::Result {
+                    specificity: *specificity,
+                    get_content: WrappedGetContentSourceContentVc::new(
+                        *get_content,
+                        self.processor,
+                    )
+                    .into(),
+                }
+                .cell()
+            }
+        })
+    }
+}
+
+/// A WrappedGetContentSourceContent simply wraps the get_content of a
+/// [ContentSourceResult], allowing us to process whatever
+/// [ContentSourceContent] it would have returned.
+///
+/// This is the second of 2 steps, implementing the processing of
+/// ContentSourceContent. The first step in [WrappedContentSource] handles
+/// ContentSourceResult.
+#[turbo_tasks::value]
+struct WrappedGetContentSourceContent {
+    inner: GetContentSourceContentVc,
+    processor: ContentSourceProcessorVc,
+}
+
+#[turbo_tasks::value_impl]
+impl WrappedGetContentSourceContentVc {
+    #[turbo_tasks::function]
+    fn new(inner: GetContentSourceContentVc, processor: ContentSourceProcessorVc) -> Self {
+        WrappedGetContentSourceContent { inner, processor }.cell()
+    }
+}
+
+#[turbo_tasks::value_impl]
+impl GetContentSourceContent for WrappedGetContentSourceContent {
+    #[turbo_tasks::function]
+    fn vary(&self) -> ContentSourceDataVaryVc {
+        self.inner.vary()
+    }
+
+    #[turbo_tasks::function]
+    async fn get(&self, data: Value<ContentSourceData>) -> Result<ContentSourceContentVc> {
+        let res = self.inner.get(data);
+        Ok(self.processor.process(res))
+    }
+}