Fix bug in coroutine scheduling. Add ByteStream implementation and si…

…mplify Python symbol visitor. (#1633)

* Fix bug in coroutine scheduling. Add ByteStream implementation and simplify Python symbol visitor.

* Bug in coroutine scheduling:

In the previous iteration, the runtime crate aws-smithy-http-server-python
was exposing the python application implementation as a struct, which
was wrapped by the codegenerated App to allow to dynamically building
the router.

This caused scheduling of coroutines (handlers prefixed with async def)
to block becuse we were passing the Python eventloop of the parent
process that was stored pre-fork().

This commit changes the runtime PyApp to become a trait, allowing us to
dynamically build the router post-fork() and with the right event loop.

This change also allowed us to remove a bunch of unnecessary Arc(s).

* Add ByteStream implementation

Implementation of a ByteStream type for Python that can be roughly used like this:

let b = await ByteStream.from_path("file.txt")
async for chunk in b:
    print(chunk)

Implement futures_core::stream::Stream for Python ByteStream wrapper.

The BytesStream implementation in python can now use a sync Mutex from
parking_lot because we are now using pyo3_asyncio::tokio::local_future_into_py()
to read a chunk, which supports !Send futures.

This allows to simply forward the implementation of Stream (poll_next()
and size_hint()) directly to our inner SDK ByteStream.

* Simplify Python symbol visitor

Inherit from Symbol visitor and override just what is needed to swap Python complex
types.

Signed-off-by: Bigo <[email protected]>

codegen-server-test/python/model/pokemon.smithy

@@ -13,7 +13,7 @@ use aws.protocols#restJson1
 service PokemonService {
     version: "2021-12-01",
     resources: [PokemonSpecies],
-    operations: [GetServerStatistics, EmptyOperation, HealthCheckOperation],
+    operations: [GetServerStatistics, EmptyOperation, HealthCheckOperation, StreamPokemonRadioOperation],
 }
 
 /// A Pokémon species forms the basis for at least one Pokémon.
@@ -70,6 +70,22 @@ structure GetServerStatisticsOutput {
     calls_count: Long,
 }
 
+/// Fetch the radio song from the database and stream it back as a playable audio.
+@readonly
+@http(uri: "/radio", method: "GET")
+operation StreamPokemonRadioOperation {
+    output: StreamPokemonRadioOutput,
+}
+
+@output
+structure StreamPokemonRadioOutput {
+    @httpPayload
+    data: StreamingBlob,
+}
+
+@streaming
+blob StreamingBlob
+
 list FlavorTextEntries {
     member: FlavorText
 }
@@ -127,8 +143,7 @@ structure EmptyOperationOutput { }
 /// Not yet a deep check
 @readonly
 @http(uri: "/ping", method: "GET")
-operation HealthCheckOperation {
-}
+operation HealthCheckOperation { }
 
 @error("client")
 @httpError(404)

codegen-server/python/src/main/kotlin/software/amazon/smithy/rust/codegen/server/python/smithy/PythonCodegenServerPlugin.kt

@@ -12,13 +12,12 @@ import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.shapes.ServiceShape
 import software.amazon.smithy.rust.codegen.rustlang.RustReservedWordSymbolProvider
 import software.amazon.smithy.rust.codegen.server.python.smithy.customizations.DECORATORS
-import software.amazon.smithy.rust.codegen.server.python.smithy.generators.PythonServerSymbolProvider
+import software.amazon.smithy.rust.codegen.server.python.smithy.generators.PythonServerSymbolVisitor
+import software.amazon.smithy.rust.codegen.server.python.smithy.generators.PythonStreamingShapeMetadataProvider
 import software.amazon.smithy.rust.codegen.server.smithy.customizations.ServerRequiredCustomizations
 import software.amazon.smithy.rust.codegen.smithy.BaseSymbolMetadataProvider
 import software.amazon.smithy.rust.codegen.smithy.EventStreamSymbolProvider
 import software.amazon.smithy.rust.codegen.smithy.ServerCodegenContext
-import software.amazon.smithy.rust.codegen.smithy.StreamingShapeMetadataProvider
-import software.amazon.smithy.rust.codegen.smithy.StreamingShapeSymbolProvider
 import software.amazon.smithy.rust.codegen.smithy.SymbolVisitor
 import software.amazon.smithy.rust.codegen.smithy.SymbolVisitorConfig
 import software.amazon.smithy.rust.codegen.smithy.customize.CombinedCodegenDecorator
@@ -68,20 +67,17 @@ class PythonCodegenServerPlugin : SmithyBuildPlugin {
             serviceShape: ServiceShape,
             symbolVisitorConfig: SymbolVisitorConfig,
         ) =
-            SymbolVisitor(model, serviceShape = serviceShape, config = symbolVisitorConfig)
+            // Rename a set of symbols that do not implement `PyClass` and have been wrapped in
+            // `aws_smithy_http_server_python::types`.
+            PythonServerSymbolVisitor(model, serviceShape = serviceShape, config = symbolVisitorConfig)
                 // Generate different types for EventStream shapes (e.g. transcribe streaming)
                 .let {
                     EventStreamSymbolProvider(symbolVisitorConfig.runtimeConfig, it, model, CodegenTarget.SERVER)
                 }
-                // Generate [ByteStream] instead of `Blob` for streaming binary shapes (e.g. S3 GetObject)
-                .let { StreamingShapeSymbolProvider(it, model) }
-                // Rename a set of symbols that do not implement `PyClass` and have been wrapped in
-                // `aws_smithy_http_server_python::types`.
-                .let { PythonServerSymbolProvider(it, model) }
                 // Add Rust attributes (like `#[derive(PartialEq)]`) to generated shapes
                 .let { BaseSymbolMetadataProvider(it, model, additionalAttributes = listOf()) }
                 // Streaming shapes need different derives (e.g. they cannot derive Eq)
-                .let { StreamingShapeMetadataProvider(it, model) }
+                .let { PythonStreamingShapeMetadataProvider(it, model) }
                 // Rename shapes that clash with Rust reserved words & and other SDK specific features e.g. `send()` cannot
                 // be the name of an operation input
                 .let { RustReservedWordSymbolProvider(it, model) }

codegen-server/python/src/main/kotlin/software/amazon/smithy/rust/codegen/server/python/smithy/PythonServerCargoDependency.kt

@@ -23,6 +23,7 @@ object PythonServerCargoDependency {
     val TowerHttp: CargoDependency = CargoDependency("tower-http", CratesIo("0.3"), features = setOf("trace"))
     val Hyper: CargoDependency = CargoDependency("hyper", CratesIo("0.14.12"), features = setOf("server", "http1", "http2", "tcp", "stream"))
     val NumCpus: CargoDependency = CargoDependency("num_cpus", CratesIo("1.13"))
+    val ParkingLot: CargoDependency = CargoDependency("parking_lot", CratesIo("0.12"))
 
     fun SmithyHttpServer(runtimeConfig: RuntimeConfig) = runtimeConfig.runtimeCrate("http-server")
     fun SmithyHttpServerPython(runtimeConfig: RuntimeConfig) = runtimeConfig.runtimeCrate("http-server-python")

codegen-server/python/src/main/kotlin/software/amazon/smithy/rust/codegen/server/python/smithy/PythonServerCodegenVisitor.kt

@@ -27,7 +27,6 @@ import software.amazon.smithy.rust.codegen.smithy.generators.BuilderGenerator
 import software.amazon.smithy.rust.codegen.smithy.generators.CodegenTarget
 import software.amazon.smithy.rust.codegen.smithy.generators.implBlock
 import software.amazon.smithy.rust.codegen.util.getTrait
-import software.amazon.smithy.rust.codegen.util.hasStreamingMember
 
 /**
  * Entrypoint for Python server-side code generation. This class will walk the in-memory model and
@@ -83,9 +82,6 @@ class PythonServerCodegenVisitor(
      * This function _does not_ generate any serializers.
      */
     override fun structureShape(shape: StructureShape) {
-        if (shape.hasStreamingMember(model)) {
-            throw CodegenException("Streaming members are not supported in Python yet")
-        }
         logger.info("[python-server-codegen] Generating a structure $shape")
         rustCrate.useShapeWriter(shape) { writer ->
             // Use Python specific structure generator that adds the #[pyclass] attribute

codegen-server/python/src/main/kotlin/software/amazon/smithy/rust/codegen/server/python/smithy/PythonServerRuntimeType.kt

@@ -21,6 +21,9 @@ object PythonServerRuntimeType {
     fun Blob(runtimeConfig: RuntimeConfig) =
         RuntimeType("Blob", PythonServerCargoDependency.SmithyHttpServerPython(runtimeConfig), "${runtimeConfig.crateSrcPrefix}_http_server_python::types")
 
+    fun ByteStream(runtimeConfig: RuntimeConfig) =
+        RuntimeType("ByteStream", PythonServerCargoDependency.SmithyHttpServerPython(runtimeConfig), "${runtimeConfig.crateSrcPrefix}_http_server_python::types")
+
     fun DateTime(runtimeConfig: RuntimeConfig) =
         RuntimeType("DateTime", PythonServerCargoDependency.SmithyHttpServerPython(runtimeConfig), "${runtimeConfig.crateSrcPrefix}_http_server_python::types")
 

codegen-server/python/src/main/kotlin/software/amazon/smithy/rust/codegen/server/python/smithy/PythonServerSymbolProvider.kt

@@ -8,102 +8,107 @@ package software.amazon.smithy.rust.codegen.server.python.smithy.generators
 import software.amazon.smithy.codegen.core.Symbol
 import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.shapes.BlobShape
-import software.amazon.smithy.model.shapes.ListShape
-import software.amazon.smithy.model.shapes.MapShape
 import software.amazon.smithy.model.shapes.MemberShape
-import software.amazon.smithy.model.shapes.SetShape
+import software.amazon.smithy.model.shapes.ServiceShape
 import software.amazon.smithy.model.shapes.Shape
-import software.amazon.smithy.model.shapes.ShapeVisitor
+import software.amazon.smithy.model.shapes.StringShape
 import software.amazon.smithy.model.shapes.StructureShape
 import software.amazon.smithy.model.shapes.TimestampShape
 import software.amazon.smithy.model.shapes.UnionShape
-import software.amazon.smithy.model.traits.EnumDefinition
-import software.amazon.smithy.rust.codegen.rustlang.RustType
+import software.amazon.smithy.rust.codegen.rustlang.RustMetadata
 import software.amazon.smithy.rust.codegen.server.python.smithy.PythonServerRuntimeType
-import software.amazon.smithy.rust.codegen.smithy.MaybeRenamed
+import software.amazon.smithy.rust.codegen.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.smithy.RustSymbolProvider
+import software.amazon.smithy.rust.codegen.smithy.SymbolMetadataProvider
+import software.amazon.smithy.rust.codegen.smithy.SymbolVisitor
 import software.amazon.smithy.rust.codegen.smithy.SymbolVisitorConfig
-import software.amazon.smithy.rust.codegen.smithy.rustType
+import software.amazon.smithy.rust.codegen.smithy.expectRustMetadata
 import software.amazon.smithy.rust.codegen.smithy.shape
-import software.amazon.smithy.rust.codegen.util.toPascalCase
-import software.amazon.smithy.rust.codegen.util.toSnakeCase
+import software.amazon.smithy.rust.codegen.smithy.traits.SyntheticInputTrait
+import software.amazon.smithy.rust.codegen.smithy.traits.SyntheticOutputTrait
+import software.amazon.smithy.rust.codegen.util.hasStreamingMember
+import software.amazon.smithy.rust.codegen.util.hasTrait
+import software.amazon.smithy.rust.codegen.util.isStreaming
 
-// Wrapping symbol visitor provider allowing to implement symbol providers that can recursively replace symbols in nested shapes.
-open class PythonWrappingVisitingSymbolProvider(private val base: RustSymbolProvider, private val model: Model) : ShapeVisitor.Default<Symbol>(), RustSymbolProvider {
-    override fun getDefault(shape: Shape): Symbol {
-        return base.toSymbol(shape)
-    }
-
-    override fun config(): SymbolVisitorConfig {
-        return base.config()
-    }
-
-    override fun toEnumVariantName(definition: EnumDefinition): MaybeRenamed? {
-        return base.toEnumVariantName(definition)
-    }
+/**
+ * Symbol visitor  allowing that recursively replace symbols in nested shapes.
+ *
+ * Input / output / error structures can refer to complex types like the ones implemented inside
+ * `aws_smithy_types` (a good example is `aws_smithy_types::Blob`).
+ * `aws_smithy_http_server_python::types` wraps those types that do not implement directly the
+ * `pyo3::PyClass` trait and cannot be shared safely with Python, providing an idiomatic Python / Rust API.
+ *
+ * This symbol provider ensures types not implementing `pyo3::PyClass` are swapped with their wrappers from
+ * `aws_smithy_http_server_python::types`.
+ */
+class PythonServerSymbolVisitor(
+    private val model: Model,
+    private val serviceShape: ServiceShape?,
+    private val config: SymbolVisitorConfig,
+) : SymbolVisitor(model, serviceShape, config) {
+    private val runtimeConfig = config().runtimeConfig
 
-    override fun toMemberName(shape: MemberShape): String = when (val container = model.expectShape(shape.container)) {
-        is StructureShape -> shape.memberName.toSnakeCase()
-        is UnionShape -> shape.memberName.toPascalCase()
-        else -> error("unexpected container shape: $container")
-    }
+    override fun toSymbol(shape: Shape): Symbol {
+        val initial = shape.accept(this)
 
-    override fun listShape(shape: ListShape): Symbol {
-        val inner = toSymbol(shape.member)
-        return symbolBuilder(RustType.Vec(inner.rustType())).addReference(inner).build()
-    }
+        if (shape !is MemberShape) {
+            return initial
+        }
+        val target = model.expectShape(shape.target)
+        val container = model.expectShape(shape.container)
 
-    override fun mapShape(shape: MapShape): Symbol {
-        val target = model.expectShape(shape.key.target)
-        require(target.isStringShape) { "unexpected key shape: ${shape.key}: $target [keys must be strings]" }
-        val key = toSymbol(shape.key)
-        val value = toSymbol(shape.value)
-        return symbolBuilder(RustType.HashMap(key.rustType(), value.rustType())).addReference(key)
-            .addReference(value).build()
-    }
+        // We are only targetting non syntetic inputs and outputs.
+        if (!container.hasTrait<SyntheticOutputTrait>() && !container.hasTrait<SyntheticInputTrait>()) {
+            return initial
+        }
 
-    override fun setShape(shape: SetShape): Symbol {
-        val inner = toSymbol(shape.member)
-        val builder = if (model.expectShape(shape.member.target).isStringShape) {
-            symbolBuilder(RustType.HashSet(inner.rustType()))
+        // We are only targeting streaming blobs as the rest of the symbols do not change if streaming is enabled.
+        // For example a TimestampShape doesn't become a different symbol when streaming is involved, but BlobShape
+        // become a ByteStream.
+        return if (target is BlobShape && shape.isStreaming(model)) {
+            PythonServerRuntimeType.ByteStream(config().runtimeConfig).toSymbol()
         } else {
-            // only strings get put into actual sets because floats are unhashable
-            symbolBuilder(RustType.Vec(inner.rustType()))
+            initial
         }
-        return builder.addReference(inner).build()
     }
 
-    override fun memberShape(shape: MemberShape): Symbol {
-        return toSymbol(model.expectShape(shape.target))
-    }
-
-    override fun toSymbol(shape: Shape): Symbol {
-        return shape.accept(this)
+    override fun timestampShape(shape: TimestampShape?): Symbol {
+        return PythonServerRuntimeType.DateTime(runtimeConfig).toSymbol()
     }
 
-    private fun symbolBuilder(rustType: RustType): Symbol.Builder {
-        return Symbol.builder().rustType(rustType)
-            .definitionFile("python.rs")
+    override fun blobShape(shape: BlobShape?): Symbol {
+        return PythonServerRuntimeType.Blob(runtimeConfig).toSymbol()
     }
 }
 
 /**
- * Input / output / error structures can refer to complex types like the ones implemented inside
- * `aws_smithy_types` (a good example is `aws_smithy_types::Blob`).
- * `aws_smithy_http_server_python::types` wraps those types that do not implement directly the
- * `pyo3::PyClass` trait and cannot be shared safely with Python, providing an idiomatic Python / Rust API.
+ * SymbolProvider to drop the PartialEq bounds in streaming shapes
  *
- * This symbol provider ensures types not implementing `pyo3::PyClass` are swapped with their wrappers from
- * `aws_smithy_http_server_python::types`.
+ * Streaming shapes equality cannot be checked without reading the body. Because of this, these shapes
+ * do not implement `PartialEq`.
+ *
+ * Note that since streaming members can only be used on the root shape, this can only impact input and output shapes.
  */
-class PythonServerSymbolProvider(private val base: RustSymbolProvider, private val model: Model) : PythonWrappingVisitingSymbolProvider(base, model) {
-    private val runtimeConfig = config().runtimeConfig
+class PythonStreamingShapeMetadataProvider(private val base: RustSymbolProvider, private val model: Model) : SymbolMetadataProvider(base) {
+    override fun memberMeta(memberShape: MemberShape): RustMetadata {
+        return base.toSymbol(memberShape).expectRustMetadata()
+    }
 
-    override fun timestampShape(shape: TimestampShape?): Symbol {
-        return PythonServerRuntimeType.DateTime(runtimeConfig).toSymbol()
+    override fun structureMeta(structureShape: StructureShape): RustMetadata {
+        val baseMetadata = base.toSymbol(structureShape).expectRustMetadata()
+        return if (structureShape.hasStreamingMember(model)) {
+            baseMetadata.withoutDerives(RuntimeType.PartialEq)
+        } else baseMetadata
     }
 
-    override fun blobShape(shape: BlobShape?): Symbol {
-        return PythonServerRuntimeType.Blob(runtimeConfig).toSymbol()
+    override fun unionMeta(unionShape: UnionShape): RustMetadata {
+        val baseMetadata = base.toSymbol(unionShape).expectRustMetadata()
+        return if (unionShape.hasStreamingMember(model)) {
+            baseMetadata.withoutDerives(RuntimeType.PartialEq)
+        } else baseMetadata
+    }
+
+    override fun enumMeta(stringShape: StringShape): RustMetadata {
+        return base.toSymbol(stringShape).expectRustMetadata()
     }
 }