From 17f741301db0bb08da0eafe8a338e5efd8a4b5df Mon Sep 17 00:00:00 2001
From: Ryo Yamashita <qryxip@gmail.com>
Date: Mon, 25 Nov 2024 23:47:07 +0900
Subject: [PATCH] fix: cherry pick the fix for `download-binaries` (#10)

Fixes the problem where building multiple ort-sys `download-binaries`
fail.

Cherry picked from eb516468, eb516468, and pykeio/ort#323.

Refs: VOICEVOX/voicevox_core#859, pykeio/ort#322
Refs: eb51646860bc2fc5728d879bdf542c0537acc6ea
Refs: effd7a76e627fdb5382d3092fb2b1f183d4ae90a, pykeio/ort#323
Co-authored-by: Carson M. <carson@pyke.io>
---
 ort-sys/build.rs            | 39 +++++++++++++++++++++++++++----------
 ort-sys/src/internal/mod.rs | 16 +++++++++++++--
 ort-sys/src/lib.rs          |  3 +++
 3 files changed, 46 insertions(+), 12 deletions(-)

diff --git a/ort-sys/build.rs b/ort-sys/build.rs
index dc0672c7..23e94a8d 100644
--- a/ort-sys/build.rs
+++ b/ort-sys/build.rs
@@ -1,5 +1,5 @@
 use std::{
-	env, fs,
+	env, fs, io,
 	path::{Path, PathBuf},
 	process::Command
 };
@@ -12,9 +12,9 @@ const ORT_ENV_SYSTEM_LIB_PROFILE: &str = "ORT_LIB_PROFILE";
 
 const DIST_TABLE: &str = include_str!("dist.txt");
 
-#[path = "src/internal/dirs.rs"]
-mod dirs;
-use self::dirs::cache_dir;
+#[path = "src/internal/mod.rs"]
+mod internal;
+use self::internal::dirs::cache_dir;
 
 #[cfg(feature = "download-binaries")]
 fn fetch_file(source_url: &str) -> Vec<u8> {
@@ -416,21 +416,40 @@ fn prepare_libort_dir() -> (PathBuf, bool) {
 
 			let (prebuilt_url, prebuilt_hash) = dist.unwrap();
 
-			let mut cache_dir = cache_dir()
+			let bin_extract_dir = cache_dir()
 				.expect("could not determine cache directory")
 				.join("dfbin")
 				.join(target)
 				.join(prebuilt_hash);
-			if fs::create_dir_all(&cache_dir).is_err() {
-				cache_dir = env::var("OUT_DIR").unwrap().into();
-			}
 
 			let ort_extract_dir = prebuilt_url.split('/').last().unwrap().strip_suffix(".tgz").unwrap();
-			let lib_dir = cache_dir.join(ort_extract_dir);
+			let lib_dir = bin_extract_dir.join(ort_extract_dir);
 			if !lib_dir.exists() {
 				let downloaded_file = fetch_file(prebuilt_url);
 				assert!(verify_file(&downloaded_file, prebuilt_hash), "hash of downloaded ONNX Runtime binary does not match!");
-				extract_tgz(&downloaded_file, &cache_dir);
+
+				let mut temp_extract_dir = bin_extract_dir
+					.parent()
+					.unwrap()
+					.join(format!("tmp.{}_{prebuilt_hash}", self::internal::random_identifier()));
+				let mut should_rename = true;
+				if fs::create_dir_all(&temp_extract_dir).is_err() {
+					temp_extract_dir = env::var("OUT_DIR").unwrap().into();
+					should_rename = false;
+				}
+				extract_tgz(&downloaded_file, &temp_extract_dir);
+				if should_rename {
+					match std::fs::rename(&temp_extract_dir, &bin_extract_dir) {
+						Ok(()) => {}
+						Err(e) => {
+							if bin_extract_dir.exists() {
+								let _ = fs::remove_dir_all(temp_extract_dir);
+							} else {
+								panic!("failed to extract downloaded binaries: {e}");
+							}
+						}
+					}
+				}
 			}
 
 			static_link_prerequisites(true);
diff --git a/ort-sys/src/internal/mod.rs b/ort-sys/src/internal/mod.rs
index 7516f3c7..6cb610ab 100644
--- a/ort-sys/src/internal/mod.rs
+++ b/ort-sys/src/internal/mod.rs
@@ -1,4 +1,16 @@
+use std::hash::{BuildHasher, Hasher, RandomState};
+
 pub mod dirs;
 
-#[cfg(feature = "download-binaries")]
-include!(concat!(env!("OUT_DIR"), "/downloaded_version.rs"));
+pub fn random_identifier() -> String {
+	let mut state = RandomState::new().build_hasher().finish();
+	std::iter::repeat_with(move || {
+		state ^= state << 13;
+		state ^= state >> 7;
+		state ^= state << 17;
+		state
+	})
+	.take(12)
+	.map(|i| b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"[i as usize % 62] as char)
+	.collect()
+}
diff --git a/ort-sys/src/lib.rs b/ort-sys/src/lib.rs
index ee08e634..660e76ff 100644
--- a/ort-sys/src/lib.rs
+++ b/ort-sys/src/lib.rs
@@ -10,6 +10,9 @@
 #[doc(hidden)]
 pub mod internal;
 
+#[cfg(feature = "download-binaries")]
+include!(concat!(env!("OUT_DIR"), "/downloaded_version.rs"));
+
 pub const ORT_API_VERSION: u32 = 17;
 
 pub use std::ffi::{c_char, c_int, c_ulong, c_ulonglong, c_ushort, c_void};