write JobContext to file

oceanbase · Sep 24, 2024 · 9bee724 · 9bee724
1 parent b880623
commit 9bee724
Show file tree

Hide file tree

Showing 8 changed files with 119 additions and 8 deletions.
diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/task/caller/JobCallerBuilder.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/task/caller/JobCallerBuilder.java
@@ -15,8 +15,13 @@
  */
 package com.oceanbase.odc.service.task.caller;
 
+import java.io.File;
+import java.nio.charset.Charset;
 import java.util.Map;
 
+import org.apache.commons.io.FileUtils;
+
+import com.oceanbase.odc.common.util.StringUtils;
 import com.oceanbase.odc.service.task.config.JobConfigurationHolder;
 import com.oceanbase.odc.service.task.config.TaskFrameworkProperties;
 import com.oceanbase.odc.service.task.constants.JobEnvKeyConstants;
@@ -36,7 +41,24 @@ public class JobCallerBuilder {
     public static JobCaller buildProcessCaller(JobContext context) {
         Map<String, String> environments = new JobEnvironmentFactory().build(context, TaskRunMode.PROCESS);
         JobUtils.encryptEnvironments(environments);
-
+        /**
+         * write JobContext to file in case of exceeding the environments size limit; set the file path in
+         * the environment instead
+         */
+        String jobContextFilePath = JobUtils.getExecutorDataPath() + "/" + StringUtils.uuid() + ".enc";
+        try {
+            FileUtils.writeStringToFile(new File(jobContextFilePath),
+                    environments.get(JobEnvKeyConstants.ODC_JOB_CONTEXT),
+                    Charset.defaultCharset());
+        } catch (Exception ex) {
+            FileUtils.deleteQuietly(new File(jobContextFilePath));
+            throw new RuntimeException("Failed to write job context to file: " + jobContextFilePath, ex);
+        }
+        environments.put(JobEnvKeyConstants.ODC_JOB_CONTEXT_FILE_PATH,
+                JobUtils.encrypt(environments.get(JobEnvKeyConstants.ENCRYPT_KEY),
+                        environments.get(JobEnvKeyConstants.ENCRYPT_SALT), jobContextFilePath));
+        // remove JobContext from environments
+        environments.remove(JobEnvKeyConstants.ODC_JOB_CONTEXT);
         ProcessConfig config = new ProcessConfig();
         config.setEnvironments(environments);
 

diff --git a/...-service/src/main/java/com/oceanbase/odc/service/task/caller/JobEnvironmentEncryptor.java b/...-service/src/main/java/com/oceanbase/odc/service/task/caller/JobEnvironmentEncryptor.java
@@ -42,6 +42,7 @@ public class JobEnvironmentEncryptor {
             JobEnvKeyConstants.ODC_EXECUTOR_DATABASE_PASSWORD,
             JobEnvKeyConstants.ODC_OBJECT_STORAGE_CONFIGURATION,
             JobEnvKeyConstants.ODC_PROPERTY_ENCRYPTION_SALT,
+            JobEnvKeyConstants.ODC_JOB_CONTEXT_FILE_PATH,
             JobEnvKeyConstants.ODC_JOB_CONTEXT);
 
     private final AtomicBoolean encrypted = new AtomicBoolean(false);
@@ -85,5 +86,13 @@ public void decrypt(@NonNull Map<String, String> environments) {
         });
     }
 
+    public String encrypt(String key, String salt, String raw) {
+        TextEncryptor textEncryptor = Encryptors.aesBase64(key, salt);
+        return textEncryptor.encrypt(raw);
+    }
 
+    public String decrypt(String key, String salt, String encrypted) {
+        TextEncryptor textEncryptor = Encryptors.aesBase64(key, salt);
+        return textEncryptor.decrypt(encrypted);
+    }
 }
diff --git a/...dc-service/src/main/java/com/oceanbase/odc/service/task/constants/JobEnvKeyConstants.java b/...dc-service/src/main/java/com/oceanbase/odc/service/task/constants/JobEnvKeyConstants.java
@@ -27,12 +27,16 @@ public class JobEnvKeyConstants {
 
     public static final String ODC_JOB_CONTEXT = "ODC_JOB_CONTEXT";
 
+    public static final String ODC_JOB_CONTEXT_FILE_PATH = "ODC_JOB_CONTEXT_FILE_PATH";
+
     public static final String ODC_TASK_RUN_MODE = "ODC_TASK_RUN_MODE";
 
     public static final String ODC_BOOT_MODE = "ODC_BOOT_MODE";
 
     public static final String ODC_LOG_DIRECTORY = "odc.log.directory";
 
+    public static final String ODC_DATA_DIRECTORY = "file.storage.dir";
+
     public static final String ODC_EXECUTOR_PORT = "ODC_EXECUTOR_PORT";
 
     public static final String ODC_SERVICE_HOST = "ODC_SERVICE_HOST";

diff --git a/...er/odc-service/src/main/java/com/oceanbase/odc/service/task/executor/TaskApplication.java b/...er/odc-service/src/main/java/com/oceanbase/odc/service/task/executor/TaskApplication.java
@@ -24,6 +24,7 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.core.LoggerContext;
+import org.eclipse.jgit.util.StringUtils;
 
 import com.oceanbase.odc.common.trace.TaskContextHolder;
 import com.oceanbase.odc.common.trace.TraceContextHolder;
@@ -90,7 +91,8 @@ private void init(String[] args) {
         log.info("decrypt environment variables success.");
 
         // 3 step: get JobContext from environment
-        context = JobContextProviderFactory.create().provide();
+        context = JobContextProviderFactory.create(SystemUtils.getEnvOrProperty(JobEnvKeyConstants.ODC_TASK_RUN_MODE))
+                .provide();
         log.info("initial job context success.");
 
         // 4 step: trace taskId in log4j2 context
@@ -151,9 +153,14 @@ private void setLog4JConfigXml() {
     }
 
     private void validEnvValues() {
-        validNotBlank(JobEnvKeyConstants.ODC_JOB_CONTEXT);
-        validNotBlank(JobEnvKeyConstants.ODC_BOOT_MODE);
         validNotBlank(JobEnvKeyConstants.ODC_TASK_RUN_MODE);
+        if (StringUtils.equalsIgnoreCase("PROCESS",
+                SystemUtils.getEnvOrProperty(JobEnvKeyConstants.ODC_TASK_RUN_MODE))) {
+            validNotBlank(JobEnvKeyConstants.ODC_JOB_CONTEXT_FILE_PATH);
+        } else {
+            validNotBlank(JobEnvKeyConstants.ODC_JOB_CONTEXT);
+        }
+        validNotBlank(JobEnvKeyConstants.ODC_BOOT_MODE);
         validNotBlank(JobEnvKeyConstants.ENCRYPT_SALT);
         validNotBlank(JobEnvKeyConstants.ENCRYPT_KEY);
         validNotBlank(JobEnvKeyConstants.ODC_EXECUTOR_USER_ID);

diff --git a/.../main/java/com/oceanbase/odc/service/task/executor/context/JobContextProviderFactory.java b/.../main/java/com/oceanbase/odc/service/task/executor/context/JobContextProviderFactory.java
@@ -16,14 +16,23 @@
 
 package com.oceanbase.odc.service.task.executor.context;
 
+import com.oceanbase.odc.service.task.enums.TaskRunMode;
+
+import lombok.NonNull;
+
 /**
  * @author gaoda.xy
  * @date 2023/11/23 13:55
  */
 public class JobContextProviderFactory {
 
-    public static JobContextProvider create() {
-        return new DefaultJobContextProvider();
+    public static JobContextProvider create(@NonNull String taskRunMode) {
+        if (taskRunMode.equalsIgnoreCase(TaskRunMode.PROCESS.name())) {
+            return new ProcessJobContextProvider();
+        } else if (taskRunMode.equalsIgnoreCase(TaskRunMode.K8S.name())) {
+            return new K8sJobContextProvider();
+        } else {
+            throw new RuntimeException("Unsupported task run mode: " + taskRunMode);
+        }
     }
-
 }
diff --git a/...or/context/DefaultJobContextProvider.java → ...ecutor/context/K8sJobContextProvider.java b/...or/context/DefaultJobContextProvider.java → ...ecutor/context/K8sJobContextProvider.java
@@ -25,7 +25,7 @@
  * @author gaoda.xy
  * @date 2023/11/22 20:21
  */
-public class DefaultJobContextProvider implements JobContextProvider {
+public class K8sJobContextProvider implements JobContextProvider {
 
     @Override
     public JobContext provide() {

diff --git a/.../main/java/com/oceanbase/odc/service/task/executor/context/ProcessJobContextProvider.java b/.../main/java/com/oceanbase/odc/service/task/executor/context/ProcessJobContextProvider.java
@@ -0,0 +1,52 @@
+/*
+ * Copyright (c) 2023 OceanBase.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.oceanbase.odc.service.task.executor.context;
+
+import java.io.File;
+import java.nio.charset.Charset;
+
+import org.apache.commons.io.FileUtils;
+
+import com.oceanbase.odc.common.json.JsonUtils;
+import com.oceanbase.odc.common.util.SystemUtils;
+import com.oceanbase.odc.service.task.caller.DefaultJobContext;
+import com.oceanbase.odc.service.task.caller.JobContext;
+import com.oceanbase.odc.service.task.constants.JobEnvKeyConstants;
+import com.oceanbase.odc.service.task.util.JobUtils;
+
+/**
+ * @Author: Lebie
+ * @Date: 2024/9/24 17:27
+ * @Description: []
+ */
+public class ProcessJobContextProvider implements JobContextProvider {
+    @Override
+    public JobContext provide() {
+        String encryptedJobContextJson;
+        try {
+            encryptedJobContextJson = FileUtils
+                    .readFileToString(new File(System.getProperty(JobEnvKeyConstants.ODC_JOB_CONTEXT_FILE_PATH)),
+                            Charset.defaultCharset());
+        } catch (Exception ex) {
+            throw new RuntimeException("read job context file failed, ex=", ex);
+        } finally {
+            FileUtils.deleteQuietly(new File(System.getProperty(JobEnvKeyConstants.ODC_JOB_CONTEXT_FILE_PATH)));
+        }
+        String rawJobContextJson = JobUtils.decrypt(SystemUtils.getEnvOrProperty(JobEnvKeyConstants.ENCRYPT_KEY),
+                SystemUtils.getEnvOrProperty(JobEnvKeyConstants.ENCRYPT_SALT), encryptedJobContextJson);
+        return JsonUtils.fromJson(rawJobContextJson, DefaultJobContext.class);
+    }
+}
diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/task/util/JobUtils.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/task/util/JobUtils.java
@@ -180,4 +180,12 @@ public static void putEnvToSysProperties(String environmentKey) {
     public static void encryptEnvironments(Map<String, String> environments) {
         new JobEnvironmentEncryptor().encrypt(environments);
     }
+
+    public static String encrypt(String key, String salt, String raw) {
+        return new JobEnvironmentEncryptor().encrypt(key, salt, raw);
+    }
+
+    public static String decrypt(String key, String salt, String encrypted) {
+        return new JobEnvironmentEncryptor().decrypt(key, salt, encrypted);
+    }
 }