Ensure that task not executed are returned.

4 files changed, 16 insertions, 8 deletions

diff --git a/storage/src/vespa/storage/persistence/asynchandler.cpp b/storage/src/vespa/storage/persistence/asynchandler.cpp
index b8ed6b8ec91..7f8da8e76e7 100644
--- a/storage/src/vespa/storage/persistence/asynchandler.cpp
+++ b/storage/src/vespa/storage/persistence/asynchandler.cpp
@@ -104,7 +104,7 @@ AsyncHandler::handleRunTask(RunTaskCommand& cmd, MessageTracker::UP tracker) con
     });
     spi::Bucket bucket(cmd.getBucket());
     auto onDone = std::make_unique<ResultTaskOperationDone>(_sequencedExecutor, cmd.getBucketId(), std::move(task));
-    cmd.task().run(bucket, std::make_shared<vespalib::KeepAlive<decltype(onDone)>>(std::move(onDone)));
+    cmd.run(bucket, std::make_shared<vespalib::KeepAlive<decltype(onDone)>>(std::move(onDone)));
     return tracker;
 }
 
diff --git a/storage/src/vespa/storage/persistence/filestorage/filestormanager.cpp b/storage/src/vespa/storage/persistence/filestorage/filestormanager.cpp
index a46b4205570..16b06bb822e 100644
--- a/storage/src/vespa/storage/persistence/filestorage/filestormanager.cpp
+++ b/storage/src/vespa/storage/persistence/filestorage/filestormanager.cpp
@@ -978,7 +978,10 @@ FileStorManager::execute(const spi::Bucket &bucket, std::unique_ptr<spi::BucketT
     StorBucketDatabase::WrappedEntry entry(_component.getBucketDatabase(bucket.getBucketSpace()).get(
             bucket.getBucketId(), "FileStorManager::execute"));
     if (entry.exist()) {
-        _filestorHandler->schedule(std::make_shared<RunTaskCommand>(bucket, std::move(task)));
+        auto cmd = std::make_shared<RunTaskCommand>(bucket, std::move(task));
+        if ( ! _filestorHandler->schedule(cmd) ) {
+            task = cmd->stealTask();
+        }
     }
     return task;
 }
diff --git a/storage/src/vespa/storage/persistence/messages.cpp b/storage/src/vespa/storage/persistence/messages.cpp
index 7ccb3ee895d..beae22a429b 100644
--- a/storage/src/vespa/storage/persistence/messages.cpp
+++ b/storage/src/vespa/storage/persistence/messages.cpp
@@ -187,9 +187,7 @@ RunTaskCommand::RunTaskCommand(const spi::Bucket &bucket, std::unique_ptr<spi::B
     : api::InternalCommand(ID),
       _task(std::move(task)),
       _bucket(bucket)
-{
-    assert(_task);
-}
+{ }
 
 RunTaskCommand::~RunTaskCommand() = default;
 
@@ -203,6 +201,14 @@ RunTaskCommand::print(std::ostream& out, bool verbose, const std::string& indent
     }
 }
 
+void
+RunTaskCommand::run(const spi::Bucket & bucket, std::shared_ptr<vespalib::IDestructorCallback> onComplete)
+{
+    if (_task) {
+        _task->run(bucket, std::move(onComplete));
+    }
+}
+
 RunTaskReply::RunTaskReply(const RunTaskCommand& cmd)
     : api::InternalReply(ID, cmd)
 {}
diff --git a/storage/src/vespa/storage/persistence/messages.h b/storage/src/vespa/storage/persistence/messages.h
index 043747d10d2..50834782f39 100644
--- a/storage/src/vespa/storage/persistence/messages.h
+++ b/storage/src/vespa/storage/persistence/messages.h
@@ -249,11 +249,10 @@ public:
 
     document::Bucket getBucket() const override { return _bucket.getBucket(); }
     std::unique_ptr<api::StorageReply> makeReply() override;
-    spi::BucketTask & task() & {
-        return *_task;
-    }
+    void run(const spi::Bucket & bucket, std::shared_ptr<vespalib::IDestructorCallback> onComplete);
 
     void print(std::ostream& out, bool verbose, const std::string& indent) const override;
+    std::unique_ptr<spi::BucketTask> stealTask() { return std::move(_task); }
 private:
     std::unique_ptr<spi::BucketTask> _task;
     spi::Bucket                      _bucket;