feat(outbound-pg): Reuse connection during request lifecycle

Also, pg_backend_pid endpoint into outbound-pg example.

Before the fix:

$ curl -i localhost:3000/pg_backend_pid
HTTP/1.1 500 Internal Server Error
content-length: 0
date: Tue, 27 Sep 2022 14:03:50 GMT

thread '<unnamed>' panicked at 'assertion failed: `(left == right)`
  left: `592913`,
  right: `592914`', src/lib.rs:112:5

After the fix:

$ curl -i localhost:3000/pg_backend_pid
HTTP/1.1 200 OK
content-length: 23
date: Tue, 27 Sep 2022 14:07:14 GMT

pg_backend_pid: 595194

Ideally, this fix has to be covered by an integration test rather than
manual testing through examples, but the testing environment has to be
set up first.

Refs: #667.

Signed-off-by: Konstantin Shabanov <mail@etehtsea.me>

Author: etehtsea

crates/outbound-pg/src/lib.rs

@@ -1,9 +1,9 @@
 use anyhow::anyhow;
 use outbound_pg::*;
+use std::collections::HashMap;
 use tokio_postgres::{
-    tls::NoTlsStream,
     types::{ToSql, Type},
-    Connection, NoTls, Row, Socket,
+    Client, NoTls, Row,
 };
 
 pub use outbound_pg::add_to_linker;
@@ -16,8 +16,9 @@ use wit_bindgen_wasmtime::{async_trait, wasmtime::Linker};
 wit_bindgen_wasmtime::export!({paths: ["../../wit/ephemeral/outbound-pg.wit"], async: *});
 
 /// A simple implementation to support outbound pg connection
-#[derive(Default, Clone)]
-pub struct OutboundPg;
+pub struct OutboundPg {
+    pub connections: HashMap<String, Client>,
+}
 
 impl HostComponent for OutboundPg {
     type State = Self;
@@ -33,7 +34,9 @@ impl HostComponent for OutboundPg {
         &self,
         _component: &spin_manifest::CoreComponent,
     ) -> anyhow::Result<Self::State> {
-        Ok(Self)
+        let connections = std::collections::HashMap::new();
+
+        Ok(Self { connections })
     }
 }
 
@@ -45,19 +48,16 @@ impl outbound_pg::OutboundPg for OutboundPg {
         statement: &str,
         params: Vec<ParameterValue<'_>>,
     ) -> Result<u64, PgError> {
-        let (client, connection) = tokio_postgres::connect(address, NoTls)
-            .await
-            .map_err(|e| PgError::ConnectionFailed(format!("{:?}", e)))?;
-
-        spawn(connection);
-
         let params: Vec<&(dyn ToSql + Sync)> = params
             .iter()
             .map(to_sql_parameter)
             .collect::<anyhow::Result<Vec<_>>>()
             .map_err(|e| PgError::ValueConversionFailed(format!("{:?}", e)))?;
 
-        let nrow = client
+        let nrow = self
+            .get_client(address)
+            .await
+            .map_err(|e| PgError::ConnectionFailed(format!("{:?}", e)))?
             .execute(statement, params.as_slice())
             .await
             .map_err(|e| PgError::QueryFailed(format!("{:?}", e)))?;
@@ -71,19 +71,16 @@ impl outbound_pg::OutboundPg for OutboundPg {
         statement: &str,
         params: Vec<ParameterValue<'_>>,
     ) -> Result<RowSet, PgError> {
-        let (client, connection) = tokio_postgres::connect(address, NoTls)
-            .await
-            .map_err(|e| PgError::ConnectionFailed(format!("{:?}", e)))?;
-
-        spawn(connection);
-
         let params: Vec<&(dyn ToSql + Sync)> = params
             .iter()
             .map(to_sql_parameter)
             .collect::<anyhow::Result<Vec<_>>>()
             .map_err(|e| PgError::BadParameter(format!("{:?}", e)))?;
 
-        let results = client
+        let results = self
+            .get_client(address)
+            .await
+            .map_err(|e| PgError::ConnectionFailed(format!("{:?}", e)))?
             .query(statement, params.as_slice())
             .await
             .map_err(|e| PgError::QueryFailed(format!("{:?}", e)))?;
@@ -246,10 +243,26 @@ fn convert_entry(row: &Row, index: usize) -> Result<DbValue, tokio_postgres::Err
     Ok(value)
 }
 
-fn spawn(connection: Connection<Socket, NoTlsStream>) {
+impl OutboundPg {
+    async fn get_client(&mut self, address: &str) -> anyhow::Result<&Client> {
+        let client = match self.connections.entry(address.to_owned()) {
+            std::collections::hash_map::Entry::Occupied(o) => o.into_mut(),
+            std::collections::hash_map::Entry::Vacant(v) => v.insert(build_client(address).await?),
+        };
+        Ok(client)
+    }
+}
+
+async fn build_client(address: &str) -> anyhow::Result<Client> {
+    tracing::log::debug!("Build new connection: {}", address);
+
+    let (client, connection) = tokio_postgres::connect(address, NoTls).await?;
+
     tokio::spawn(async move {
         if let Err(e) = connection.await {
             tracing::warn!("Postgres connection error: {}", e);
         }
     });
+
+    Ok(client)
 }

crates/trigger/src/lib.rs

@@ -138,7 +138,9 @@ pub fn add_default_host_components<T: Default + Send + 'static>(
     builder.add_host_component(outbound_redis::OutboundRedis {
         connections: Arc::new(RwLock::new(HashMap::new())),
     })?;
-    builder.add_host_component(outbound_pg::OutboundPg)?;
+    builder.add_host_component(outbound_pg::OutboundPg {
+        connections: HashMap::new(),
+    })?;
     Ok(())
 }
 

examples/rust-outbound-pg/src/lib.rs

@@ -22,6 +22,7 @@ fn process(req: Request) -> Result<Response> {
     match req.uri().path() {
         "/read" => read(req),
         "/write" => write(req),
+        "/pg_backend_pid" => pg_backend_pid(req),
         _ => Ok(http::Response::builder()
             .status(404)
             .body(Some("Not found".into()))?),
@@ -96,6 +97,27 @@ fn write(_req: Request) -> Result<Response> {
         .body(Some(response.into()))?)
 }
 
+fn pg_backend_pid(_req: Request) -> Result<Response> {
+    let address = std::env::var(DB_URL_ENV)?;
+    let sql = "SELECT pg_backend_pid()";
+
+    let get_pid = || {
+        let rowset = pg::query(&address, sql, &[])
+            .map_err(|e| anyhow!("Error executing Postgres query: {:?}", e))?;
+
+        let row = &rowset.rows[0];
+        as_int(&row[0])
+    };
+
+    assert_eq!(get_pid()?, get_pid()?);
+
+    let response = format!("pg_backend_pid: {}\n", get_pid()?);
+
+    Ok(http::Response::builder()
+        .status(200)
+        .body(Some(response.into()))?)
+}
+
 fn as_owned_string(value: &pg::DbValue) -> anyhow::Result<String> {
     match value {
         pg::DbValue::Str(s) => Ok(s.to_owned()),