rewrite extension client and upgrade tokio (awslabs#120)

Author: bnusunny

.rustfmt.toml

@@ -1,5 +1,3 @@
 max_width=120
-unstable_features = true
-imports_granularity = "Crate"
 reorder_imports = true
 edition = "2021"

Cargo.lock

@@ -19,9 +19,8 @@ exclude = ["examples"]
 [dependencies]
 http = "0.2"
 hyper = { version = "0.14", features = ["client"] }
-lambda-extension = "0.7.0"
 lambda_http = "0.7.2"
-tokio = { version = "1.20.0", features = [
+tokio = { version = "1.24", features = [
     "macros",
     "io-util",
     "sync",

Cargo.toml

@@ -52,4 +52,4 @@ build-LambdaAdapterLayerArm64:
 	DOCKER_BUILDKIT=1 docker build --build-arg TARGET_PLATFORM=linux/arm64 --build-arg ARCH=aarch64 -o $(ARTIFACTS_DIR)/extensions .
 
 fmt:
-	cargo +nightly fmt --all
+	cargo fmt --all

Makefile

@@ -17,13 +17,13 @@ use std::{
 use encoding_rs::{Encoding, UTF_8};
 use http::{
     header::{HeaderName, HeaderValue},
-    Uri,
+    Method, StatusCode, Uri,
 };
 use hyper::{
     body::HttpBody,
     client::{Client, HttpConnector},
+    Body,
 };
-use lambda_extension::Extension;
 use lambda_http::aws_lambda_events::serde_json;
 pub use lambda_http::Error;
 use lambda_http::{Body as LambdaBody, Request, RequestExt, Response};
@@ -137,13 +137,31 @@ impl Adapter {
     pub fn register_default_extension(&self) {
         // register as an external extension
         tokio::task::spawn(async move {
-            match Extension::new().with_events(&[]).run().await {
-                Ok(_) => {}
-                Err(err) => {
-                    tracing::error!(err = err, "extension terminated unexpectedly");
-                    panic!("extension thread execution");
-                }
+            let aws_lambda_runtime_api: String =
+                env::var("AWS_LAMBDA_RUNTIME_API").unwrap_or_else(|_| "127.0.0.1:9001".to_string());
+            let client = hyper::Client::new();
+            let register_req = hyper::Request::builder()
+                .method(Method::POST)
+                .uri(format!("http://{aws_lambda_runtime_api}/2020-01-01/extension/register"))
+                .header("Lambda-Extension-Name", "lambda-adapter")
+                .body(Body::from("{ \"events\": [] }"))
+                .unwrap();
+            let register_res = client.request(register_req).await.unwrap();
+            if register_res.status() != StatusCode::OK {
+                panic!("extension registration failure");
             }
+            let next_req = hyper::Request::builder()
+                .method(Method::GET)
+                .uri(format!(
+                    "http://{aws_lambda_runtime_api}/2020-01-01/extension/event/next"
+                ))
+                .header(
+                    "Lambda-Extension-Identifier",
+                    register_res.headers().get("Lambda-Extension-Identifier").unwrap(),
+                )
+                .body(Body::empty())
+                .unwrap();
+            client.request(next_req).await.unwrap();
         });
     }
 

src/lib.rs

@@ -8,10 +8,9 @@ use httpmock::{
     Method::{DELETE, GET, POST, PUT},
     MockServer,
 };
-use lambda_extension::Service;
-
 use lambda_http::Body;
 use lambda_web_adapter::{Adapter, AdapterOptions, Protocol};
+use tower::Service;
 
 #[test]
 fn test_adapter_options_from_env() {