feat: encode/decode from preset coder's URN

sdks/rust/Cargo.toml

@@ -43,7 +43,7 @@ rand = "0.7"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 serde_yaml = "0.9.14"
-strum_macros = "0.24"
+strum = { version = "0.24", features = ["derive"] }
 tokio = { version = "1.0", features = ["rt-multi-thread", "macros", "sync", "time"] }
 tokio-util = "0.7.4"
 tokio-stream = "0.1"

sdks/rust/src/coders/register_coders/mod.rs

@@ -48,8 +48,8 @@ macro_rules! register_coders {
         }
 
         #[ctor::ctor]
-        fn init_coder_from_urn() {
-            $crate::worker::CODER_FROM_URN.set($crate::worker::CoderFromUrn {
+        fn init_custom_coder_from_urn() {
+            $crate::worker::CUSTOM_CODER_FROM_URN.set($crate::worker::CustomCoderFromUrn {
                 enc: encode_from_urn,
                 dec: decode_from_urn,
             }).unwrap();

sdks/rust/src/coders/required_coders.rs

@@ -133,23 +133,22 @@ impl fmt::Debug for BytesCoder {
 }
 
 /// A coder for a key-value pair
-#[derive(Clone, Debug)]
 pub struct KVCoder<KV> {
     phantom: PhantomData<KV>,
 }
 
 impl<K, V> CoderUrn for KVCoder<KV<K, V>>
 where
-    K: ElemType + Clone + fmt::Debug,
-    V: ElemType + Clone + fmt::Debug,
+    K: ElemType,
+    V: ElemType,
 {
     const URN: &'static str = KV_CODER_URN;
 }
 
 impl<K, V> Coder for KVCoder<KV<K, V>>
 where
-    K: ElemType + Clone + fmt::Debug,
-    V: ElemType + Clone + fmt::Debug,
+    K: ElemType,
+    V: ElemType,
 {
     /// Encode the input element (a key-value pair) into a byte output stream. They key and value are encoded one after the
     /// other (first key, then value). The key is encoded with `Context::NeedsDelimiters`, while the value is encoded with
@@ -173,10 +172,20 @@ where
     }
 }
 
+impl<K, V> fmt::Debug for KVCoder<KV<K, V>>
+where
+    K: ElemType,
+    V: ElemType,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("KVCoder").finish()
+    }
+}
+
 impl<K, V> Default for KVCoder<KV<K, V>>
 where
-    K: Send,
-    V: Send,
+    K: ElemType,
+    V: ElemType,
 {
     fn default() -> Self {
         Self {

sdks/rust/src/coders/urns.rs

@@ -16,6 +16,44 @@
  * limitations under the License.
  */
 
+use strum::{EnumDiscriminants, EnumIter};
+
+#[derive(Debug, EnumDiscriminants, EnumIter)]
+pub(crate) enum PresetCoderUrn {
+    Bytes,
+    Kv,
+    Iterable,
+    StrUtf8,
+    VarInt,
+    Unit,
+    GeneralObject,
+}
+
+impl PresetCoderUrn {
+    pub(crate) fn as_str(&self) -> &str {
+        match self {
+            // ******* Standard coders *******
+            Self::Bytes => BYTES_CODER_URN,
+            Self::Kv => KV_CODER_URN,
+            Self::Iterable => ITERABLE_CODER_URN,
+
+            // ******* Required coders *******
+            Self::StrUtf8 => STR_UTF8_CODER_URN,
+            Self::VarInt => VARINT_CODER_URN,
+
+            // ******* Rust coders *******
+            Self::Unit => UNIT_CODER_URN,
+            Self::GeneralObject => GENERAL_OBJECT_CODER_URN,
+        }
+    }
+}
+
+impl AsRef<str> for PresetCoderUrn {
+    fn as_ref(&self) -> &str {
+        self.as_str()
+    }
+}
+
 // ******* Standard coders *******
 pub const BYTES_CODER_URN: &str = "beam:coder:bytes:v1";
 pub const KV_CODER_URN: &str = "beam:coder:kvcoder:v1";

sdks/rust/src/lib.rs

@@ -28,6 +28,3 @@ pub mod proto;
 pub mod runners;
 pub mod transforms;
 pub mod worker;
-
-#[macro_use]
-extern crate strum_macros;

sdks/rust/src/worker/coder_from_urn/custom_coder_from_urn/mod.rs

@@ -0,0 +1,47 @@
+use std::fmt;
+
+use once_cell::sync::OnceCell;
+
+use crate::coders::{DecodeFromUrnFn, EncodeFromUrnFn};
+
+/// The visibility is `pub` because this is used internally from `register_coders!` macro.
+pub static CUSTOM_CODER_FROM_URN: OnceCell<CustomCoderFromUrn> = OnceCell::new();
+
+/// The visibility is `pub` because this is instantiated internally from `register_coders!` macro.
+pub struct CustomCoderFromUrn {
+    pub enc: EncodeFromUrnFn,
+    pub dec: DecodeFromUrnFn,
+}
+
+impl CustomCoderFromUrn {
+    pub(in crate::worker::coder_from_urn) fn global() -> &'static CustomCoderFromUrn {
+        CUSTOM_CODER_FROM_URN
+            .get()
+            .expect("you might forget calling `register_coders!(CustomCoder1, CustomCoder2)`")
+    }
+
+    pub(in crate::worker::coder_from_urn) fn encode_from_urn(
+        &self,
+        urn: &str,
+        elem: &dyn crate::elem_types::ElemType,
+        writer: &mut dyn std::io::Write,
+        context: &crate::coders::Context,
+    ) -> Result<usize, std::io::Error> {
+        (self.enc)(urn, elem, writer, context)
+    }
+
+    pub(in crate::worker::coder_from_urn) fn decode_from_urn(
+        &self,
+        urn: &str,
+        reader: &mut dyn std::io::Read,
+        context: &crate::coders::Context,
+    ) -> Result<Box<dyn crate::elem_types::ElemType>, std::io::Error> {
+        (self.dec)(urn, reader, context)
+    }
+}
+
+impl fmt::Debug for CustomCoderFromUrn {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("CodersFromUrn").finish()
+    }
+}

sdks/rust/src/worker/coder_from_urn/mod.rs

@@ -1,47 +1,33 @@
-use std::fmt;
+mod custom_coder_from_urn;
+pub use custom_coder_from_urn::{CustomCoderFromUrn, CUSTOM_CODER_FROM_URN};
 
-use once_cell::sync::OnceCell;
+use crate::worker::coder_from_urn::preset_coder_from_urn::PresetCoderFromUrn;
 
-use crate::coders::{DecodeFromUrnFn, EncodeFromUrnFn};
+mod preset_coder_from_urn;
 
-/// Called internally from `register_coders!` macro.
-pub static CODER_FROM_URN: OnceCell<CoderFromUrn> = OnceCell::new();
-
-/// Called internally from `register_coders!` macro.
-pub struct CoderFromUrn {
-    pub enc: EncodeFromUrnFn,
-    pub dec: DecodeFromUrnFn,
-}
+pub(in crate::worker) struct CoderFromUrn;
 
 impl CoderFromUrn {
-    pub fn global() -> &'static CoderFromUrn {
-        CODER_FROM_URN
-            .get()
-            .expect("you might forget calling `register_coders!(CustomCoder1, CustomCoder2)`")
-    }
-
-    pub fn encode_from_urn(
-        &self,
+    pub(in crate::worker) fn encode_from_urn(
         urn: &str,
         elem: &dyn crate::elem_types::ElemType,
         writer: &mut dyn std::io::Write,
         context: &crate::coders::Context,
     ) -> Result<usize, std::io::Error> {
-        (self.enc)(urn, elem, writer, context)
+        PresetCoderFromUrn::encode_from_urn(urn, elem, writer, context).unwrap_or_else(|| {
+            let custom = CustomCoderFromUrn::global();
+            (custom.enc)(urn, elem, writer, context)
+        })
     }
 
-    pub fn decode_from_urn(
-        &self,
+    pub(in crate::worker) fn decode_from_urn(
         urn: &str,
         reader: &mut dyn std::io::Read,
         context: &crate::coders::Context,
     ) -> Result<Box<dyn crate::elem_types::ElemType>, std::io::Error> {
-        (self.dec)(urn, reader, context)
-    }
-}
-
-impl fmt::Debug for CoderFromUrn {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        f.debug_struct("CodersFromUrn").finish()
+        PresetCoderFromUrn::decode_from_urn(urn, reader, context).unwrap_or_else(|| {
+            let custom = CustomCoderFromUrn::global();
+            (custom.dec)(urn, reader, context)
+        })
     }
 }

sdks/rust/src/worker/coder_from_urn/preset_coder_from_urn/mod.rs

@@ -0,0 +1,51 @@
+use crate::coders::{
+    required_coders::BytesCoder, rust_coders::GeneralObjectCoder, standard_coders::StrUtf8Coder,
+    urns::PresetCoderUrn, Coder,
+};
+use strum::IntoEnumIterator;
+
+#[derive(Eq, PartialEq, Debug)]
+pub(in crate::worker::coder_from_urn) struct PresetCoderFromUrn;
+
+impl PresetCoderFromUrn {
+    /// Returns `None` if the urn is not a preset coder urn.
+    pub(in crate::worker) fn encode_from_urn(
+        urn: &str,
+        elem: &dyn crate::elem_types::ElemType,
+        writer: &mut dyn std::io::Write,
+        context: &crate::coders::Context,
+    ) -> Option<Result<usize, std::io::Error>> {
+        let opt_variant = PresetCoderUrn::iter().find(|variant| variant.as_str() == urn);
+
+        opt_variant.map(|variant| match variant {
+            PresetCoderUrn::Bytes => BytesCoder::default().encode(elem, writer, context),
+            PresetCoderUrn::Kv => todo!("create full type including components (not only urn but also full proto maybe required"),
+            PresetCoderUrn::Iterable => todo!("create full type including components (not only urn but also full proto maybe required"),
+            PresetCoderUrn::StrUtf8 => StrUtf8Coder::default().encode(elem, writer, context),
+            PresetCoderUrn::VarInt => todo!("create full type including components (not only urn but also full proto maybe required"),
+            PresetCoderUrn::Unit => todo!("make UnitCoder"),
+            PresetCoderUrn::GeneralObject => {
+                GeneralObjectCoder::default().encode(elem, writer, context)
+            }
+        })
+    }
+
+    /// Returns `None` if the urn is not a preset coder urn.
+    pub(in crate::worker) fn decode_from_urn(
+        urn: &str,
+        reader: &mut dyn std::io::Read,
+        context: &crate::coders::Context,
+    ) -> Option<Result<Box<dyn crate::elem_types::ElemType>, std::io::Error>> {
+        let opt_variant = PresetCoderUrn::iter().find(|variant| variant.as_str() == urn);
+
+        opt_variant.map(|variant| match variant {
+            PresetCoderUrn::Bytes => BytesCoder::default().decode(reader, context),
+            PresetCoderUrn::Kv => todo!("create full type including components (not only urn but also full proto maybe required"),
+            PresetCoderUrn::Iterable => todo!("create full type including components (not only urn but also full proto maybe required"),
+            PresetCoderUrn::StrUtf8 => StrUtf8Coder::default().decode(reader, context),
+            PresetCoderUrn::VarInt => todo!("create full type including components (not only urn but also full proto maybe required"),
+            PresetCoderUrn::Unit => todo!("make UnitCoder"),
+            PresetCoderUrn::GeneralObject => GeneralObjectCoder::default().decode(reader, context),
+        })
+    }
+}