allow to configure qdrant init file with env variable (qdrant#2316)

* allow to configure qdrant init file with env variable * Return path type from get_init_file_path --------- Co-authored-by: timvisee <tim@visee.me>
AKJUS · Jul 31, 2023 · 4fd9b4d · 4fd9b4d
1 parent 52a8fb3
commit 4fd9b4d
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 7 deletions.
diff --git a/src/startup.rs b/src/startup.rs
@@ -2,12 +2,19 @@
 
 use std::backtrace::Backtrace;
 use std::panic;
+use std::path::PathBuf;
 
 use log::LevelFilter;
 
 use crate::common::error_reporting::ErrorReporter;
 
-const INITIALIZED_FILE: &str = ".qdrant-initialized";
+const DEFAULT_INITIALIZED_FILE: &str = ".qdrant-initialized";
+
+fn get_init_file_path() -> PathBuf {
+    std::env::var("QDRANT_INIT_FILE_PATH")
+        .map(PathBuf::from)
+        .unwrap_or_else(|_| DEFAULT_INITIALIZED_FILE.into())
+}
 
 pub fn setup_logger(log_level: &str) {
     let is_info = log_level.to_ascii_uppercase() == "INFO";
@@ -61,16 +68,17 @@ pub fn setup_panic_hook(reporting_enabled: bool, reporting_id: String) {
 /// Creates a file that indicates that the server has been started.
 /// This file is used to check if the server has been been successfully started before potential kill.
 pub fn touch_started_file_indicator() {
-    if let Err(err) = std::fs::write(INITIALIZED_FILE, "") {
+    if let Err(err) = std::fs::write(get_init_file_path(), "") {
         log::warn!("Failed to create init file indicator: {}", err);
     }
 }
 
 /// Removes a file that indicates that the server has been started.
 /// Use before server initialization to avoid false positives.
 pub fn remove_started_file_indicator() {
-    if std::path::Path::new(INITIALIZED_FILE).exists() {
-        if let Err(err) = std::fs::remove_file(INITIALIZED_FILE) {
+    let path = get_init_file_path();
+    if path.exists() {
+        if let Err(err) = std::fs::remove_file(path) {
             log::warn!("Failed to remove init file indicator: {}", err);
         }
     }

diff --git a/tools/entrypoint.sh b/tools/entrypoint.sh
@@ -39,14 +39,14 @@ if [ $EXIT_CODE != 137 ]; then
     exit $EXIT_CODE
 fi
 
-IS_INITIALIZED_FILE='.qdrant-initialized'
+QDRANT_INIT_FILE_PATH=${QDRANT_INIT_FILE_PATH:-'.qdrant-initialized'}
 RECOVERY_MESSAGE="Qdrant was killed during initialization. Most likely it's Out-of-Memory.
 Please check memory consumption, increase memory limit or remove some collections and restart"
 
 # Check that qdrant was initialized
-# Qdrant creates IS_INITIALIZED_FILE file after initialization
+# Qdrant creates QDRANT_INIT_FILE_PATH file after initialization
 # So if it doesn't exist, qdrant was killed during initialization
-if [ ! -f "$IS_INITIALIZED_FILE" ]; then
+if [ ! -f "$QDRANT_INIT_FILE_PATH" ]; then
     # Run qdrant in recovery mode.
     # No collection operations are allowed in recovery mode except for removing collections
     QDRANT__STORAGE__RECOVERY_MODE="$RECOVERY_MESSAGE" ./qdrant $@ &