fix: race conditions in data store save (#13372)

* fix: race conditions in data store save

* Whitespace change to kick CI

* Flsuh saves to disk

* Remove log

.changeset/empty-cloths-sort.md

@@ -0,0 +1,5 @@
+---
+'astro': patch
+---
+
+Fixes a bug that caused some very large data stores to save incomplete data.

packages/astro/src/content/content-layer.ts

@@ -296,11 +296,11 @@ export class ContentLayer {
 		);
 		await fs.mkdir(this.#settings.config.cacheDir, { recursive: true });
 		await fs.mkdir(this.#settings.dotAstroDir, { recursive: true });
-		await this.#store.writeToDisk();
 		const assetImportsFile = new URL(ASSET_IMPORTS_FILE, this.#settings.dotAstroDir);
 		await this.#store.writeAssetImports(assetImportsFile);
 		const modulesImportsFile = new URL(MODULES_IMPORTS_FILE, this.#settings.dotAstroDir);
 		await this.#store.writeModuleImports(modulesImportsFile);
+		await this.#store.waitUntilSaveComplete();
 		logger.info('Synced content');
 		if (this.#settings.config.experimental.contentIntellisense) {
 			await this.regenerateCollectionFileManifest();

packages/astro/src/content/mutable-data-store.ts

@@ -25,6 +25,9 @@ export class MutableDataStore extends ImmutableDataStore {
 	#assetsSaveTimeout: NodeJS.Timeout | undefined;
 	#modulesSaveTimeout: NodeJS.Timeout | undefined;
 
+	#savePromise: Promise<void> | undefined;
+	#savePromiseResolve: (() => void) | undefined;
+
 	#dirty = false;
 	#assetsDirty = false;
 	#modulesDirty = false;
@@ -152,15 +155,36 @@ export default new Map([\n${lines.join(',\n')}]);
 		this.#modulesDirty = false;
 	}
 
+	#maybeResolveSavePromise() {
+		if (
+			!this.#saveTimeout &&
+			!this.#assetsSaveTimeout &&
+			!this.#modulesSaveTimeout &&
+			this.#savePromiseResolve
+		) {
+			this.#savePromiseResolve();
+			this.#savePromiseResolve = undefined;
+			this.#savePromise = undefined;
+		}
+	}
+
 	#writeAssetsImportsDebounced() {
 		this.#assetsDirty = true;
 		if (this.#assetsFile) {
 			if (this.#assetsSaveTimeout) {
 				clearTimeout(this.#assetsSaveTimeout);
 			}
-			this.#assetsSaveTimeout = setTimeout(() => {
+
+			if (!this.#savePromise) {
+				this.#savePromise = new Promise<void>((resolve) => {
+					this.#savePromiseResolve = resolve;
+				});
+			}
+
+			this.#assetsSaveTimeout = setTimeout(async () => {
 				this.#assetsSaveTimeout = undefined;
-				this.writeAssetImports(this.#assetsFile!);
+				await this.writeAssetImports(this.#assetsFile!);
+				this.#maybeResolveSavePromise();
 			}, SAVE_DEBOUNCE_MS);
 		}
 	}
@@ -171,23 +195,51 @@ export default new Map([\n${lines.join(',\n')}]);
 			if (this.#modulesSaveTimeout) {
 				clearTimeout(this.#modulesSaveTimeout);
 			}
-			this.#modulesSaveTimeout = setTimeout(() => {
+
+			if (!this.#savePromise) {
+				this.#savePromise = new Promise<void>((resolve) => {
+					this.#savePromiseResolve = resolve;
+				});
+			}
+
+			this.#modulesSaveTimeout = setTimeout(async () => {
 				this.#modulesSaveTimeout = undefined;
-				this.writeModuleImports(this.#modulesFile!);
+				await this.writeModuleImports(this.#modulesFile!);
+				this.#maybeResolveSavePromise();
 			}, SAVE_DEBOUNCE_MS);
 		}
 	}
 
+	// Skips the debounce and writes to disk immediately
+	async #saveToDiskNow() {
+		if (this.#saveTimeout) {
+			clearTimeout(this.#saveTimeout);
+		}
+		this.#saveTimeout = undefined;
+		if (this.#file) {
+			await this.writeToDisk();
+		}
+		this.#maybeResolveSavePromise();
+	}
+
 	#saveToDiskDebounced() {
 		this.#dirty = true;
 		if (this.#saveTimeout) {
 			clearTimeout(this.#saveTimeout);
 		}
-		this.#saveTimeout = setTimeout(() => {
+
+		if (!this.#savePromise) {
+			this.#savePromise = new Promise<void>((resolve) => {
+				this.#savePromiseResolve = resolve;
+			});
+		}
+
+		this.#saveTimeout = setTimeout(async () => {
 			this.#saveTimeout = undefined;
 			if (this.#file) {
-				this.writeToDisk();
+				await this.writeToDisk();
 			}
+			this.#maybeResolveSavePromise();
 		}, SAVE_DEBOUNCE_MS);
 	}
 
@@ -331,6 +383,19 @@ export default new Map([\n${lines.join(',\n')}]);
 		};
 	}
 
+	/**
+	 * Returns a promise that resolves when all pending saves are complete.
+	 * This includes any in-progress debounced saves for the data store, asset imports, and module imports.
+	 */
+	async waitUntilSaveComplete(): Promise<void> {
+		// If there's no save promise, all saves are complete
+		if (!this.#savePromise) {
+			return Promise.resolve();
+		}
+		await this.#saveToDiskNow();
+		return this.#savePromise;
+	}
+
 	toString() {
 		return devalue.stringify(this._collections);
 	}
@@ -343,8 +408,9 @@ export default new Map([\n${lines.join(',\n')}]);
 			throw new AstroError(AstroErrorData.UnknownFilesystemError);
 		}
 		try {
-			await this.#writeFileAtomic(this.#file, this.toString());
+			// Mark as clean before writing to disk so that it catches any changes that happen during the write
 			this.#dirty = false;
+			await this.#writeFileAtomic(this.#file, this.toString());
 		} catch (err) {
 			throw new AstroError(AstroErrorData.UnknownFilesystemError, { cause: err });
 		}

packages/astro/src/content/runtime.ts

@@ -20,7 +20,7 @@ import {
 	unescapeHTML,
 } from '../runtime/server/index.js';
 import { CONTENT_LAYER_TYPE, IMAGE_IMPORT_PREFIX } from './consts.js';
-import { type DataEntry, type ImmutableDataStore, globalDataStore } from './data-store.js';
+import { type DataEntry, globalDataStore } from './data-store.js';
 import type { ContentLookupMap } from './utils.js';
 
 type LazyImport = () => Promise<any>;
@@ -620,10 +620,6 @@ async function render({
 }
 
 export function createReference({ lookupMap }: { lookupMap: ContentLookupMap }) {
-	// We're handling it like this to avoid needing an async schema. Not ideal, but should
-	// be safe because the store will already have been loaded by the time this is called.
-	let store: ImmutableDataStore | null = null;
-	globalDataStore.get().then((s) => (store = s));
 	return function reference(collection: string) {
 		return z
 			.union([
@@ -645,14 +641,6 @@ export function createReference({ lookupMap }: { lookupMap: ContentLookupMap })
 						| { slug: string; collection: string },
 					ctx,
 				) => {
-					if (!store) {
-						ctx.addIssue({
-							code: ZodIssueCode.custom,
-							message: `**${ctx.path.join('.')}:** Reference to ${collection} could not be resolved: store not available.\nThis is an Astro bug, so please file an issue at https://github.com/withastro/astro/issues.`,
-						});
-						return;
-					}
-
 					const flattenedErrorPath = ctx.path.join('.');
 
 					if (typeof lookup === 'object') {