1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112

use crate::artifacts::Object2Updater;
use crate::artifacts::PublicSeriesId;
use crate::groups::base_artifact_uploader::BaseArtifactUploader;
use crate::groups::ArtifactUploader2d;
use crate::groups::ArtifactUploader3d;
use crate::groups::GenericArtifactUploader;
use crate::groups::RunUploader;
use crate::PublicArtifactId;
use async_channel::Receiver;
use async_trait::async_trait;
use std::ops::Deref;
use std::ops::DerefMut;
use wasm_bindgen::prelude::wasm_bindgen;

#[async_trait]
pub trait ArtifactUploadHandle<T>: Deref<Target = T> {
    async fn wait_for_upload(&self);
}

#[macro_export]
macro_rules! task_handle_impl {
    ($sub:ident $res:ident) => {
        #[wasm_bindgen]
        #[derive(Debug, Clone)]
        pub struct $sub {
            #[wasm_bindgen(getter_with_clone)]
            pub result: $res,
            pub(crate) channel: async_channel::Receiver<()>,
        }

        #[wasm_bindgen]
        impl $sub {
            pub async fn wait_for(&self) {
                // TODO(doug): Expose error for caller
                let _unused = self.channel.recv().await;
            }
        }

        #[async_trait]
        impl ArtifactUploadHandle<$res> for $sub {
            async fn wait_for_upload(&self) {
                self.wait_for().await;
            }
        }

        impl Deref for $sub {
            type Target = $res;

            fn deref(&self) -> &Self::Target {
                &self.result
            }
        }

        impl DerefMut for $sub {
            fn deref_mut(&mut self) -> &mut Self::Target {
                &mut self.result
            }
        }

        impl TaskHandle for $sub {
            type Result = $res;

            fn new(result: Self::Result, channel: Receiver<()>) -> Self {
                Self { result, channel }
            }

            fn args(self) -> (Self::Result, Receiver<()>) {
                (self.result, self.channel)
            }
        }
    };
}
// wasm-bindgen can't handle generics, so we have to implement each task handle
// manually
task_handle_impl!(PublicArtifactIdTaskHandle PublicArtifactId);
task_handle_impl!(GenericArtifactUploaderTaskHandle GenericArtifactUploader);
task_handle_impl!(ArtifactUploader2dTaskHandle ArtifactUploader2d);
task_handle_impl!(ArtifactUploader3dTaskHandle ArtifactUploader3d);
task_handle_impl!(PublicSeriesIdTaskHandle PublicSeriesId);
task_handle_impl!(RunUploaderTaskHandle RunUploader);
task_handle_impl!(Object2UpdaterTaskHandle Object2Updater);

#[derive(Debug, Clone)]
pub(crate) struct BaseArtifactUploaderTaskHandle {
    pub result: BaseArtifactUploader,
    channel: async_channel::Receiver<()>,
}

impl TaskHandle for BaseArtifactUploaderTaskHandle {
    type Result = BaseArtifactUploader;

    fn new(result: Self::Result, channel: Receiver<()>) -> Self {
        Self { result, channel }
    }

    fn args(self) -> (Self::Result, Receiver<()>) {
        (self.result, self.channel)
    }
}

pub(crate) trait TaskHandle: Sized {
    type Result;

    fn new(result: Self::Result, channel: async_channel::Receiver<()>) -> Self;

    fn args(self) -> (Self::Result, async_channel::Receiver<()>);

    fn map_handle<T: TaskHandle>(self, f: impl FnOnce(Self::Result) -> T::Result) -> T {
        let (result, channel) = self.args();
        T::new(f(result), channel)
    }
}