From 4396f38fcc6fe8179d0acf5acacbff44383003b3 Mon Sep 17 00:00:00 2001
From: Steven Fackler <sfackler@gmail.com>
Date: Thu, 25 Jul 2019 07:21:53 -0700
Subject: [PATCH] Setup socket

---
 tokio-postgres/src/connect_socket.rs | 11 ++++++++-
 tokio-postgres/tests/test/main.rs    | 36 ----------------------------
 2 files changed, 10 insertions(+), 37 deletions(-)

diff --git a/tokio-postgres/src/connect_socket.rs b/tokio-postgres/src/connect_socket.rs
index d88edbc9..bcc00ed0 100644
--- a/tokio-postgres/src/connect_socket.rs
+++ b/tokio-postgres/src/connect_socket.rs
@@ -31,7 +31,16 @@ pub async fn connect_socket(idx: usize, config: &Config) -> Result<Socket, Error
             let mut error = None;
             for addr in addrs {
                 let new_error = match connect_timeout(TcpStream::connect(&addr), config).await {
-                    Ok(socket) => return Ok(Socket::new_tcp(socket)),
+                    Ok(socket) => {
+                        socket.set_nodelay(true).map_err(Error::connect)?;
+                        if config.keepalives {
+                            socket
+                                .set_keepalive(Some(config.keepalives_idle))
+                                .map_err(Error::connect)?;
+                        }
+
+                        return Ok(Socket::new_tcp(socket));
+                    }
                     Err(e) => e,
                 };
                 error = Some(new_error);
diff --git a/tokio-postgres/tests/test/main.rs b/tokio-postgres/tests/test/main.rs
index 411929e0..5d27a43d 100644
--- a/tokio-postgres/tests/test/main.rs
+++ b/tokio-postgres/tests/test/main.rs
@@ -141,42 +141,6 @@ async fn insert_select() {
 }
 
 /*
-#[test]
-fn insert_select() {
-    let _ = env_logger::try_init();
-    let mut runtime = Runtime::new().unwrap();
-
-    let (mut client, connection) = runtime.block_on(connect("user=postgres")).unwrap();
-    let connection = connection.map_err(|e| panic!("{}", e));
-    runtime.handle().spawn(connection).unwrap();
-
-    runtime
-        .block_on(
-            client
-                .simple_query("CREATE TEMPORARY TABLE foo (id SERIAL, name TEXT)")
-                .for_each(|_| Ok(())),
-        )
-        .unwrap();
-
-    let insert = client.prepare("INSERT INTO foo (name) VALUES ($1), ($2)");
-    let select = client.prepare("SELECT id, name FROM foo ORDER BY id");
-    let prepare = insert.join(select);
-    let (insert, select) = runtime.block_on(prepare).unwrap();
-
-    let insert = client
-        .execute(&insert, &[&"alice", &"bob"])
-        .map(|n| assert_eq!(n, 2));
-    let select = client.query(&select, &[]).collect().map(|rows| {
-        assert_eq!(rows.len(), 2);
-        assert_eq!(rows[0].get::<_, i32>(0), 1);
-        assert_eq!(rows[0].get::<_, &str>(1), "alice");
-        assert_eq!(rows[1].get::<_, i32>(0), 2);
-        assert_eq!(rows[1].get::<_, &str>(1), "bob");
-    });
-    let tests = insert.join(select);
-    runtime.block_on(tests).unwrap();
-}
-
 #[test]
 fn query_portal() {
     let _ = env_logger::try_init();