[WIP] scheduler for running operations subsequently

lzy/allocator/src/main/java/ai/lzy/allocator/alloc/DeleteSessionAction.java

@@ -46,14 +46,6 @@ protected boolean isInjectedError(Error e) {
         return e instanceof InjectedFailures.TerminateException;
     }
 
-    @Override
-    protected void notifyExpired() {
-    }
-
-    @Override
-    protected void notifyFinished() {
-    }
-
     @Override
     protected void onExpired(TransactionHandle tx) {
         throw new RuntimeException("Unexpected, sessionId: '%s', op: '%s'".formatted(sessionId, id()));

lzy/allocator/src/main/java/ai/lzy/allocator/alloc/MountDynamicDiskAction.java

@@ -7,7 +7,10 @@
 import ai.lzy.allocator.volume.VolumeManager;
 import ai.lzy.logs.LogContextKey;
 import ai.lzy.longrunning.Operation;
-import ai.lzy.longrunning.OperationRunnerBase;
+import ai.lzy.longrunning.task.OpTaskAwareAction;
+import ai.lzy.longrunning.task.OperationTask;
+import ai.lzy.longrunning.task.OperationTaskScheduler;
+import ai.lzy.longrunning.task.dao.OperationTaskDao;
 import ai.lzy.model.db.TransactionHandle;
 import ai.lzy.v1.VmAllocatorApi;
 import com.google.protobuf.Any;
@@ -17,6 +20,7 @@
 import jakarta.annotation.Nullable;
 
 import java.sql.SQLException;
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -25,7 +29,7 @@
 
 import static ai.lzy.model.db.DbHelper.withRetries;
 
-public final class MountDynamicDiskAction extends OperationRunnerBase {
+public final class MountDynamicDiskAction extends OpTaskAwareAction {
     private final AllocationContext allocationContext;
     private final VolumeManager volumeManager;
     private final MountHolderManager mountHolderManager;
@@ -42,9 +46,13 @@ public final class MountDynamicDiskAction extends OperationRunnerBase {
     private Long nextId;
     private boolean mountPodsDeleted;
 
-    public MountDynamicDiskAction(Vm vm, DynamicMount dynamicMount, AllocationContext allocationContext) {
-        super(dynamicMount.mountOperationId(), String.format("Mount %s to VM %s", dynamicMount.mountName(), vm.vmId()),
-            allocationContext.storage(), allocationContext.operationsDao(), allocationContext.executor());
+    public MountDynamicDiskAction(Vm vm, DynamicMount dynamicMount, AllocationContext allocationContext,
+                                  OperationTask operationTask, OperationTaskDao operationTaskDao,
+                                  Duration leaseDuration, OperationTaskScheduler operationTaskScheduler)
+    {
+        super(operationTask, operationTaskDao, leaseDuration, dynamicMount.mountOperationId(),
+            String.format("Mount %s to VM %s", dynamicMount.mountName(), vm.vmId()), allocationContext.storage(),
+            allocationContext.operationsDao(), allocationContext.executor(), operationTaskScheduler);
         this.dynamicMount = dynamicMount;
         this.vm = vm;
         this.allocationContext = allocationContext;

lzy/allocator/src/main/java/ai/lzy/allocator/task/MountDynamicDiskResolver.java

@@ -0,0 +1,88 @@
+package ai.lzy.allocator.task;
+
+import ai.lzy.allocator.alloc.AllocationContext;
+import ai.lzy.allocator.alloc.MountDynamicDiskAction;
+import ai.lzy.allocator.alloc.dao.DynamicMountDao;
+import ai.lzy.allocator.alloc.dao.VmDao;
+import ai.lzy.allocator.model.DynamicMount;
+import ai.lzy.allocator.model.Vm;
+import ai.lzy.longrunning.task.OperationTask;
+import ai.lzy.longrunning.task.OperationTaskScheduler;
+import ai.lzy.longrunning.task.ResolverUtils;
+import ai.lzy.longrunning.task.TypedOperationTaskResolver;
+import ai.lzy.longrunning.task.dao.OperationTaskDao;
+import ai.lzy.model.db.TransactionHandle;
+import jakarta.annotation.Nullable;
+import jakarta.inject.Singleton;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+
+import java.sql.SQLException;
+import java.time.Duration;
+
+@Singleton
+public class MountDynamicDiskResolver implements TypedOperationTaskResolver {
+    private static final Logger LOG = LogManager.getLogger(MountDynamicDiskResolver.class);
+
+    private static final String TYPE = "MOUNT";
+    public static final String VM_ID_FIELD = "vm_id";
+    public static final String DYNAMIC_MOUNT_ID_FIELD = "dynamic_mount_id";
+
+    private final VmDao vmDao;
+    private final DynamicMountDao dynamicMountDao;
+    private final AllocationContext allocationContext;
+    private final OperationTaskDao operationTaskDao;
+    private final OperationTaskScheduler taskScheduler; //todo circular dependency
+    private final Duration leaseDuration;
+
+    public MountDynamicDiskResolver(VmDao vmDao, DynamicMountDao dynamicMountDao, AllocationContext allocationContext,
+                                    OperationTaskDao operationTaskDao, OperationTaskScheduler taskScheduler,
+                                    Duration leaseDuration)
+    //todo mark duration with a qualifier
+    {
+        this.vmDao = vmDao;
+        this.dynamicMountDao = dynamicMountDao;
+        this.allocationContext = allocationContext;
+        this.operationTaskDao = operationTaskDao;
+        this.taskScheduler = taskScheduler;
+        this.leaseDuration = leaseDuration;
+    }
+
+    @Override
+    public Result resolve(OperationTask opTask, @Nullable TransactionHandle tx) throws SQLException {
+        var vmId = ResolverUtils.readString(opTask.metadata(), VM_ID_FIELD);
+        if (vmId == null) {
+            LOG.error("{} field is not present in task {} metadata", VM_ID_FIELD, opTask.id());
+            return Result.BAD_STATE;
+        }
+        var dynamicMountId = ResolverUtils.readString(opTask.metadata(), DYNAMIC_MOUNT_ID_FIELD);
+        if (dynamicMountId == null) {
+            LOG.error("{} field is not present in task {} metadata", DYNAMIC_MOUNT_ID_FIELD, opTask.id());
+            return Result.BAD_STATE;
+        }
+        var vm = vmDao.get(vmId, tx);
+        if (vm == null) {
+            LOG.error("VM {} is not present for task", vmId);
+            return Result.STALE;
+        } else if (vm.status() != Vm.Status.RUNNING) {
+            LOG.error("VM {} is in wrong status: {}", vmId, vm.status());
+            return Result.STALE;
+        }
+        var dynamicMount = dynamicMountDao.get(dynamicMountId, false, tx);
+        if (dynamicMount == null) {
+            LOG.error("Dynamic mount {} is not present for task", dynamicMountId);
+            return Result.STALE;
+        } else if (dynamicMount.state() != DynamicMount.State.PENDING) {
+            LOG.error("Dynamic mount {} is in wrong status: {}", dynamicMount.id(), dynamicMount.state());
+            return Result.STALE;
+        }
+        return Result.success(new MountDynamicDiskAction(vm, dynamicMount, allocationContext, opTask, operationTaskDao,
+            leaseDuration, taskScheduler));
+    }
+
+    @Override
+    public String type() {
+        return TYPE;
+    }
+
+}

lzy/allocator/src/main/resources/db/allocator/migrations/V10__operation_task.sql

@@ -0,0 +1,21 @@
+CREATE TYPE task_status AS ENUM ('PENDING', 'RUNNING', 'FAILED', 'FINISHED', 'STALE');
+
+CREATE TYPE task_type AS ENUM ('UNMOUNT', 'MOUNT');
+
+CREATE TABLE IF NOT EXISTS operation_task(
+    id              BIGSERIAL   NOT NULL,
+    name            TEXT        NOT NULL,
+    entity_id       TEXT        NOT NULL,
+    type            task_type   NOT NULL,
+    status          task_status NOT NULL,
+    created_at      TIMESTAMP   NOT NULL,
+    updated_at      TIMESTAMP   NOT NULL,
+    metadata        JSONB       NOT NULL,
+    operation_id    TEXT,
+    worker_id       TEXT,
+    lease_till      TIMESTAMP,
+    PRIMARY KEY (id),
+    FOREIGN KEY (operation_id) REFERENCES operation(id)
+);
+
+CREATE INDEX IF NOT EXISTS task_status_entity_id_idx ON operation_task(status, entity_id, id);

lzy/long-running/pom.xml

@@ -45,6 +45,16 @@
             <artifactId>junit</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>io.zonky.test</groupId>
+            <artifactId>embedded-postgres</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-core</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
     <build>

lzy/long-running/src/main/java/ai/lzy/longrunning/OperationRunnerBase.java

@@ -32,6 +32,7 @@ public abstract class OperationRunnerBase extends ContextAwareTask {
     private final OperationDao operationsDao;
     private final OperationsExecutor executor;
     private Operation op;
+    private volatile boolean failed = false;
 
     protected OperationRunnerBase(String id, String descr, Storage storage, OperationDao operationsDao,
                                   OperationsExecutor executor)
@@ -55,6 +56,7 @@ protected final void execute() {
             }
 
             for (var step : steps()) {
+                beforeStep();
                 final var stepResult = step.get();
                 switch (stepResult.code()) {
                     case ALREADY_DONE -> { }
@@ -83,6 +85,7 @@ protected final void execute() {
                 }
             }
         } catch (Throwable e) {
+            setFailed();
             notifyFinished();
             if (e instanceof Error err && isInjectedError(err)) {
                 log.error("{} Terminated by InjectedFailure exception: {}", logPrefix, e.getMessage());
@@ -98,6 +101,18 @@ protected final void execute() {
         }
     }
 
+    protected void setFailed() {
+        failed = true;
+    }
+
+    protected boolean isFailed() {
+        return failed;
+    }
+
+    protected void beforeStep() {
+
+    }
+
     protected Map<String, String> prepareLogContext() {
         var ctx = super.prepareLogContext();
         ctx.put(LogContextKey.OPERATION_ID, id);
@@ -276,6 +291,7 @@ protected void notifyFinished() {
     }
 
     protected final void failOperation(Status status, @Nullable TransactionHandle tx) throws SQLException {
+        setFailed();
         operationsDao.fail(id, toProto(status), tx);
     }
 

lzy/long-running/src/main/java/ai/lzy/longrunning/task/DispatchingOperationTaskResolver.java

@@ -0,0 +1,56 @@
+package ai.lzy.longrunning.task;
+
+import ai.lzy.model.db.TransactionHandle;
+import com.google.common.annotations.VisibleForTesting;
+import jakarta.annotation.Nullable;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+
+import java.sql.SQLException;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+public class DispatchingOperationTaskResolver implements OperationTaskResolver {
+
+    private static final Logger LOG = LogManager.getLogger(DispatchingOperationTaskResolver.class);
+    private final Map<String, TypedOperationTaskResolver> resolvers;
+
+    public DispatchingOperationTaskResolver(List<TypedOperationTaskResolver> resolvers) {
+        this.resolvers = generateResolversMap(resolvers);
+    }
+
+    private static Map<String, TypedOperationTaskResolver> generateResolversMap(
+        List<TypedOperationTaskResolver> resolvers)
+    {
+        var types = new HashSet<String>();
+        resolvers.forEach(r -> {
+            if (!types.add(r.type())) {
+                throw new IllegalStateException("Duplicate resolver for type " + r.type());
+            }
+        });
+        return resolvers.stream()
+            .collect(Collectors.toMap(TypedOperationTaskResolver::type, r -> r));
+    }
+
+    @VisibleForTesting
+    void addResolver(TypedOperationTaskResolver resolver) {
+        resolvers.put(resolver.type(), resolver);
+    }
+
+    @Override
+    public Result resolve(OperationTask operationTask, @Nullable TransactionHandle tx) throws SQLException {
+        var resolver = resolvers.get(operationTask.type());
+        if (resolver == null) {
+            LOG.error("No resolver for task type {}. Task: {}", operationTask.type(), operationTask);
+            return Result.UNKNOWN_TASK;
+        }
+        try {
+            return resolver.resolve(operationTask, tx);
+        } catch (Exception e) {
+            LOG.error("Error while resolving task {}", operationTask.id(), e);
+            return Result.resolveError(e);
+        }
+    }
+}

lzy/long-running/src/main/java/ai/lzy/longrunning/task/OpTaskAwareAction.java

@@ -0,0 +1,77 @@
+package ai.lzy.longrunning.task;
+
+import ai.lzy.longrunning.OperationRunnerBase;
+import ai.lzy.longrunning.OperationsExecutor;
+import ai.lzy.longrunning.dao.OperationDao;
+import ai.lzy.longrunning.task.dao.OperationTaskDao;
+import ai.lzy.model.db.Storage;
+
+import java.time.Duration;
+import java.util.Map;
+
+import static ai.lzy.model.db.DbHelper.withRetries;
+
+public abstract class OpTaskAwareAction extends OperationRunnerBase {
+    private final OperationTaskScheduler operationTaskScheduler;
+    private final OperationTaskDao operationTaskDao;
+    private final Duration leaseDuration;
+    private OperationTask operationTask;
+
+    public OpTaskAwareAction(OperationTask operationTask, OperationTaskDao operationTaskDao, Duration leaseDuration,
+                             String opId, String desc, Storage storage, OperationDao operationsDao,
+                             OperationsExecutor executor, OperationTaskScheduler operationTaskScheduler)
+    {
+        super(opId, desc, storage, operationsDao, executor);
+        this.operationTask = operationTask;
+        this.operationTaskDao = operationTaskDao;
+        this.leaseDuration = leaseDuration;
+        this.operationTaskScheduler = operationTaskScheduler;
+    }
+
+    @Override
+    protected Map<String, String> prepareLogContext() {
+        var ctx = super.prepareLogContext();
+        ctx.put("task_id", String.valueOf(operationTask.id()));
+        ctx.put("task_type", operationTask.type());
+        ctx.put("task_name", operationTask.name());
+        ctx.put("task_entity_id", operationTask.entityId());
+        return ctx;
+    }
+
+    protected OperationTask task() {
+        return operationTask;
+    }
+
+    public void setTask(OperationTask operationTask) {
+        this.operationTask = operationTask;
+    }
+
+    @Override
+    protected void beforeStep() {
+        super.beforeStep();
+        try {
+            operationTask = withRetries(log(), () -> operationTaskDao.updateLease(operationTask.id(), leaseDuration,
+                null));
+        } catch (Exception e) {
+            log().error("{} Couldn't update lease on task {}", logPrefix(), task().id());
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    protected void notifyFinished() {
+        var builder = OperationTask.Update.builder();
+        if (isFailed()) {
+            builder.status(OperationTask.Status.FAILED);
+        } else {
+            builder.status(OperationTask.Status.FINISHED);
+        }
+        try {
+            operationTask = withRetries(log(), () -> operationTaskDao.update(operationTask.id(), builder.build(),
+                null));
+        } catch (Exception e) {
+            log().error("{} Couldn't finish operation task {}", logPrefix(), task().id());
+        }
+        operationTaskScheduler.releaseTask(task());
+    }
+}