Refactor Task Executor Registration to Primarily Use Configurations f…

…or MachineDefinition. (#639)

mantis-runtime-loader/src/main/java/io/mantisrx/runtime/loader/config/WorkerConfiguration.java

@@ -155,6 +155,18 @@ default Time getHeartbeatInterval() {
     @DefaultNull
     File getLocalStorageDir();
 
+    @Config("mantis.taskexecutor.hardware.cpu-cores")
+    @DefaultNull
+    Double getCpuCores();
+
+    @Config("mantis.taskexecutor.hardware.memory-in-mb")
+    @DefaultNull
+    Double getMemoryInMB();
+
+    @Config("mantis.taskexecutor.hardware.disk-in-mb")
+    @DefaultNull
+    Double getDiskInMB();
+
     @Config("mantis.taskexecutor.hardware.network-bandwidth-in-mb")
     @Default(value = "128.0")
     double getNetworkBandwidthInMB();

mantis-runtime-loader/src/main/java/io/mantisrx/runtime/loader/config/WorkerConfigurationUtils.java

@@ -53,6 +53,9 @@ public static <T extends WorkerConfiguration> WorkerConfigurationWritable toWrit
             .metricsPort(configSource.getMetricsPort())
             .metricsPublisher(configSource.getMetricsPublisher())
             .metricsPublisherFrequencyInSeconds(configSource.getMetricsPublisherFrequencyInSeconds())
+            .cpuCores(configSource.getCpuCores())
+            .memoryInMB(configSource.getMemoryInMB())
+            .diskInMB(configSource.getDiskInMB())
             .networkBandwidthInMB(configSource.getNetworkBandwidthInMB())
             .sinkPort(configSource.getSinkPort())
             .taskExecutorId(configSource.getTaskExecutorId())

mantis-runtime-loader/src/main/java/io/mantisrx/runtime/loader/config/WorkerConfigurationWritable.java

@@ -65,6 +65,9 @@ public class WorkerConfigurationWritable implements WorkerConfiguration {
 
     URI blobStoreArtifactDir;
     File localStorageDir;
+    Double cpuCores;
+    Double memoryInMB;
+    Double diskInMB;
     double networkBandwidthInMB;
     String taskExecutorAttributesStr;
     int asyncHttpClientMaxConnectionsPerHost;
@@ -243,6 +246,21 @@ public File getLocalStorageDir() {
         return this.localStorageDir;
     }
 
+    @Override
+    public Double getCpuCores() {
+        return this.cpuCores;
+    }
+
+    @Override
+    public Double getMemoryInMB() {
+        return this.memoryInMB;
+    }
+
+    @Override
+    public Double getDiskInMB() {
+        return this.diskInMB;
+    }
+
     @Override
     public double getNetworkBandwidthInMB() {
         return this.networkBandwidthInMB;

mantis-server/mantis-server-agent/src/main/java/io/mantisrx/server/agent/MachineDefinitionUtils.java

@@ -17,8 +17,17 @@
 
 import io.mantisrx.common.WorkerPorts;
 import io.mantisrx.runtime.MachineDefinition;
+import io.mantisrx.runtime.loader.config.WorkerConfiguration;
+import java.util.Optional;
 
 public class MachineDefinitionUtils {
+    /**
+     * Creates a MachineDefinition object using system parameters.
+     *
+     * @param workerPorts Worker ports object.
+     * @param networkBandwidthInMB Network bandwidth capacity in MB.
+     * @return MachineDefinition object.
+     */
     public static MachineDefinition sys(WorkerPorts workerPorts, double networkBandwidthInMB) {
         return new MachineDefinition(
             Hardware.getNumberCPUCores(),
@@ -27,4 +36,34 @@ public static MachineDefinition sys(WorkerPorts workerPorts, double networkBandw
             Hardware.getSizeOfDisk() / 1024.0 / 1024.0,
             workerPorts.getNumberOfPorts());
     }
+
+    /**
+     * Creates a MachineDefinition object from the worker configuration, falling back to system parameters.
+     *
+     * @param workerConfiguration Worker configuration object.
+     * @param workerPorts Worker ports object.
+     * @return MachineDefinition object.
+     */
+    public static MachineDefinition from(WorkerConfiguration workerConfiguration, WorkerPorts workerPorts) {
+        return fromWorkerConfiguration(workerConfiguration, workerPorts).orElseGet(() -> sys(workerPorts, workerConfiguration.getNetworkBandwidthInMB()));
+    }
+
+    /**
+     * Attempts to create a MachineDefinition from the WorkerConfiguration if all necessary parameters are present.
+     *
+     * @param workerConfiguration WorkerConfiguration object.
+     * @param workerPorts Worker ports object.
+     * @return An Optional of MachineDefinition if all parameters are present in WorkerConfiguration, else empty Optional.
+     */
+    private static Optional<MachineDefinition> fromWorkerConfiguration(WorkerConfiguration workerConfiguration, WorkerPorts workerPorts) {
+        if (workerConfiguration.getCpuCores() != null && workerConfiguration.getMemoryInMB() != null && workerConfiguration.getDiskInMB() != null) {
+            return Optional.of(new MachineDefinition(
+                workerConfiguration.getCpuCores(),
+                workerConfiguration.getMemoryInMB(),
+                workerConfiguration.getNetworkBandwidthInMB(),
+                workerConfiguration.getDiskInMB(),
+                workerPorts.getNumberOfPorts()));
+        }
+        return Optional.empty();
+    }
 }

mantis-server/mantis-server-agent/src/main/java/io/mantisrx/server/agent/TaskExecutor.java

@@ -157,8 +157,7 @@ public TaskExecutor(
             workerConfiguration.getDebugPort(), workerConfiguration.getConsolePort(),
             workerConfiguration.getCustomPort(),
             workerConfiguration.getSinkPort());
-        MachineDefinition machineDefinition =
-            MachineDefinitionUtils.sys(workerPorts, workerConfiguration.getNetworkBandwidthInMB());
+        MachineDefinition machineDefinition = MachineDefinitionUtils.from(workerConfiguration, workerPorts);
         String hostName = workerConfiguration.getExternalAddress();
 
         this.taskExecutorRegistration =