tasks: complete background tasks implementation, add integration test

Author: nikitavbv

apps/fx-test-app/src/lib.rs

@@ -13,7 +13,7 @@ use {
         sleep,
         HttpRequest,
         HttpResponse,
-        io::{http::{fetch, HttpBody}, blob::BlobGetError, kv::{self, KvSetNxPxError}, env},
+        io::{http::{fetch, HttpBody}, blob::BlobGetError, kv::{self, KvSetNxPxError}, env, tasks},
         StatusCode,
         utils::{axum::handle_request, migrations::{Migrations, Migration, SqlMigrationError}},
         random,
@@ -88,6 +88,8 @@ pub async fn http(mut req: HttpRequest) -> HttpResponse {
             .route("/test/kv/distributed-lock", get(kv_distributed_lock))
             .route("/test/kv/pubsub/subscribe", get(kv_pubsub_subscribe))
             .route("/test/kv/pubsub/publish", post(kv_pubsub_publish))
+            .route("/test/tasks/background/start", post(kv_tasks_background_start))
+            .route("/test/tasks/background/status", get(kv_tasks_background_status))
             .route("/_fx/cron", get(handle_cron))
             .route("/", get(home))
             .layer(Extension(Metrics::new())),
@@ -567,6 +569,22 @@ async fn kv_pubsub_publish(Json(req): Json<KvPubsubPublishRequest>) -> &'static
     "ok.\n"
 }
 
+async fn kv_tasks_background_start() -> &'static str {
+    tasks::run_in_background(async {
+        sleep(Duration::from_secs(1)).await;
+        kv::Kv::new("test-namespace").set("background_task_status", "done").await;
+    });
+
+    "ok.\n"
+}
+
+async fn kv_tasks_background_status() -> (StatusCode, &'static str) {
+    match kv::Kv::new("test-namespace").get("background_task_status").await {
+        Some(_) => (StatusCode::OK, "done."),
+        None => (StatusCode::NOT_FOUND, "not done yet.")
+    }
+}
+
 #[derive(Clone)]
 struct Metrics {
     test_counter: Counter,

fx-runtime/tests/integration_test.rs

@@ -700,6 +700,28 @@ async fn kv_pubsub() {
     assert_eq!("result: 129", result.text().await.unwrap());
 }
 
+#[tokio::test]
+async fn task_background() {
+    let client = init_fx_server().await;
+
+    // before task is started, kv key does not exist
+    let status = client.get("/test/tasks/background/status").send().await.unwrap();
+    assert!(status.status().as_u16() == 404);
+
+    // start task
+    let task = client.post("/test/tasks/background/start").send().await.unwrap();
+    assert!(task.status().is_success());
+
+    // after task is started, kv key does not exist immediately
+    let status = client.get("/test/tasks/background/status").send().await.unwrap();
+    assert!(status.status().as_u16() == 404);
+
+    // but after some time, it should be present
+    sleep(Duration::from_secs(2)).await;
+    let status = client.get("/test/tasks/background/status").send().await.unwrap();
+    assert!(status.status().is_success());
+}
+
 pub struct TestClient {
     client: reqwest::Client,
     base_url: String,

fx-sdk/src/api/tasks.rs

@@ -1,10 +1,10 @@
 use {
     futures::FutureExt,
-    crate::sys::FunctionResource,
+    crate::sys::{FunctionResource, add_function_resource, fx_tasks_background_spawn},
 };
 
 /// runs a task in background after request processing is finished
 pub fn run_in_background<F>(future: F) where F: Future<Output = ()> + 'static {
-    let task_resource = FunctionResource::BackgroundTask(future.boxed_local());
-    todo!()
+    let resource_id = add_function_resource(FunctionResource::BackgroundTask(future.boxed_local())).as_u64();
+    unsafe { fx_tasks_background_spawn(resource_id) };
 }

fx-sdk/src/lib.rs

@@ -35,7 +35,7 @@ use {
 };
 
 pub mod io {
-    pub use crate::api::{blob, http, kv, env};
+    pub use crate::api::{blob, http, kv, env, tasks};
 }
 
 pub mod sys;

fx-sdk/src/sys/mod.rs

@@ -47,7 +47,7 @@ pub extern "C" fn _fx_future_poll(future_resource_id: u64) -> i64 {
         match &mut *future_resource {
             FunctionResource::FunctionResponseFuture(v) => v
                 .poll_unpin(&mut context)
-                .map(|v| FunctionResource::from(v)),
+                .map(|v| Some(FunctionResource::from(v))),
             FunctionResource::HttpBody(v) => {
                 let body = std::mem::replace(&mut v.0, HttpBodyInner::Empty);
 
@@ -84,15 +84,20 @@ pub extern "C" fn _fx_future_poll(future_resource_id: u64) -> i64 {
 
                 v.0 = body;
 
-                poll
+                poll.map(|v| Some(v))
+            },
+            FunctionResource::BackgroundTask(v) => match v.poll_unpin(&mut context) {
+                std::task::Poll::Pending => std::task::Poll::Pending,
+                std::task::Poll::Ready(_) => std::task::Poll::Ready(None),
             },
             _other => panic!("resource is not future"),
         }
     });
 
     (match future_poll_result {
         Poll::Pending => FuturePollResult::Pending,
-        Poll::Ready(resource) => {
+        Poll::Ready(None) => FuturePollResult::Ready,
+        Poll::Ready(Some(resource)) => {
             replace_function_resource(&future_resource_id, resource);
             FuturePollResult::Ready
         }

fx-sdk/src/sys/resource.rs

@@ -279,7 +279,6 @@ pub fn replace_function_resource_with<F: FnOnce(FunctionResource) -> FunctionRes
     })
 }
 
-
 /// returns length of serialized resource
 pub fn serialize_function_resource(resource_id: &FunctionResourceId) -> u64 {
     FUNCTION_RESOURCES.with_borrow_mut(|resources| {