Repository: scylladb/scylla-rust-driver
Pull request #163 (2021-02-18 01:31:09)

Make ConnectionKeeper keep the connection open
## Description
This PR makes `ConnectionKeeper` detect when the connection breaks and reopen it in a loop.
Reopening is on an 8 second cooldown to not spam the network with reconnects.

Additionally `ConnectionKeeper` is now able to detect that a source port collision has occured and will retry on a different port.
First a random port is chosen, then we add `nr_shards` and wrap around until we check all possible ports (just like piodul described in #149).

## Fixed issues
Fixes: #149 

## Pre-review checklist

- [x] I have split my patch into logically separate commits.
- [x] All commit messages clearly explain what they change and why.
- [x] I added relevant tests for new features and bug fixes.
- [x] All commits compile, pass static checks and pass test.
- [x] PR description sums up the changes and reasons why they should be introduced.
- [x] I added appropriate `Fixes:` annotations to PR description.


cvybhu: @piodul @psarna please review
cvybhu: `USE KEYSPACE` from #160 was merged so I need to modify `ConnectionKeeper` to support reconnecting while still using the correct keyspace.

Patch:
diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml
index 40065df6..2510dfbf 100644
--- a/.github/workflows/rust.yml
+++ b/.github/workflows/rust.yml
@@ -11,23 +11,16 @@ env:
 
 jobs:
   build:
+
     runs-on: ubuntu-latest
-    services:
-      scylladb:
-        image: scylladb/scylla
-        ports:
-          - 9042:9042
-        options: --health-cmd "cqlsh --debug" --health-interval 5s --health-retries 10
-        volumes:
-          - ${{ github.workspace }}:/workspace
+
     steps:
     - uses: actions/checkout@v2
-    - name: Run tests
-      # test threads must be one because else database tests will run in parallel and will result in flaky tests
-      run: cargo test --verbose -- --test-threads=1
     - name: Format check
       run: cargo fmt --verbose --all -- --check
     - name: Clippy check
       run: cargo clippy --verbose --examples --tests -- -D warnings
     - name: Build
       run: cargo build --verbose --examples
+    - name: Run tests
+      run: cargo test --verbose
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index d3f88b1f..881b1c81 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -30,15 +30,22 @@ Currently, we require new PRs to compile without warnings, pass `cargo fmt` and
 
 ## Testing
 
-The easiest way to setup a running Scylla instance is to use the [Scylla Docker image](https://hub.docker.com/r/scylladb/scylla/):
-You need a running Scylla instance for the tests. 
+Tests are split into two categories: those which don't need a running Scylla instance to run, and those that do. Tests from the second category should be marked with `#[ignored]` attribute - our CI doesn't start Scylla during the test phase, so they need to be marked as ignored for the CI check to pass. This will be fixed in the future.
 
-Execute the commands below to run the tests:
+To run tests which don't need the database, use `cargo test`.
+
+To run all tests, use:
+
+```bash
+SCYLLA_URI=<ip and port of your Scylla instance> cargo test -- --ignored --test-threads=1
+```
+
+Until the CI is fixed, __please__ make sure to run all tests before submitting the PR. The easiest way to setup a running Scylla instance is to use the [Scylla Docker image](https://hub.docker.com/r/scylladb/scylla/):
 
 ```bash
 # Downloads and runs Scylla in Docker
 docker run --name scylla-ci -d scylladb/scylla
 
 # Run all tests
-SCYLLA_URI="$(docker inspect --format='{{ .NetworkSettings.IPAddress }}' scylla-ci):19042" cargo test -- --test-threads=1
+SCYLLA_URI="$(docker inspect --format='{{ .NetworkSettings.IPAddress }}' scylla-ci):19042" cargo test -- --ignored --test-threads=1
 ```
diff --git a/scylla/src/frame/value.rs b/scylla/src/frame/value.rs
index 69eec3ba..a40e54ad 100644
--- a/scylla/src/frame/value.rs
+++ b/scylla/src/frame/value.rs
@@ -14,15 +14,6 @@ pub trait Value {
 #[error("Value too big to be sent in a request - max 2GiB allowed")]
 pub struct ValueTooBig;
 
-/// Represents an unset value
-pub struct Unset;
-
-/// Enum providing a way to represent a value that might be unset
-pub enum MaybeUnset<V: Value> {
-    Unset,
-    Set(V),
-}
-
 /// Keeps a buffer with serialized Values
 /// Allows adding new Values and iterating over serialized ones
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
@@ -227,23 +218,6 @@ impl<T: Value> Value for Option<T> {
     }
 }
 
-impl Value for Unset {
-    fn serialize(&self, buf: &mut Vec<u8>) -> Result<(), ValueTooBig> {
-        // Unset serializes itself to empty value with length = -2
-        buf.put_i32(-2);
-        Ok(())
-    }
-}
-
-impl<V: Value> Value for MaybeUnset<V> {
-    fn serialize(&self, buf: &mut Vec<u8>) -> Result<(), ValueTooBig> {
-        match self {
-            MaybeUnset::Set(v) => v.serialize(buf),
-            MaybeUnset::Unset => Unset.serialize(buf),
-        }
-    }
-}
-
 // Every &impl Value should also implement Value
 impl<T: Value> Value for &T {
     fn serialize(&self, buf: &mut Vec<u8>) -> Result<(), ValueTooBig> {
@@ -320,9 +294,9 @@ impl_value_for_tuple!(T0, T1, T2, T3, T4, T5; 0, 1, 2, 3, 4, 5);
 impl_value_for_tuple!(T0, T1, T2, T3, T4, T5, T6; 0, 1, 2, 3, 4, 5, 6);
 impl_value_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7; 0, 1, 2, 3, 4, 5, 6, 7);
 impl_value_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8; 0, 1, 2, 3, 4, 5, 6, 7, 8);
-impl_value_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9;
+impl_value_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9; 
                            0, 1, 2, 3, 4, 5, 6, 7, 8, 9);
-impl_value_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10;
+impl_value_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10; 
                            0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10);
 impl_value_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11;
                            0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11);
@@ -414,9 +388,9 @@ impl_value_list_for_tuple!(T0, T1, T2, T3, T4, T5; 0, 1, 2, 3, 4, 5);
 impl_value_list_for_tuple!(T0, T1, T2, T3, T4, T5, T6; 0, 1, 2, 3, 4, 5, 6);
 impl_value_list_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7; 0, 1, 2, 3, 4, 5, 6, 7);
 impl_value_list_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8; 0, 1, 2, 3, 4, 5, 6, 7, 8);
-impl_value_list_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9;
+impl_value_list_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9; 
                            0, 1, 2, 3, 4, 5, 6, 7, 8, 9);
-impl_value_list_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10;
+impl_value_list_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10; 
                            0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10);
 impl_value_list_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11;
                            0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11);
@@ -546,17 +520,17 @@ impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5; 0, 1, 2, 3, 4, 5; 6);
 impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6; 0, 1, 2, 3, 4, 5, 6; 7);
 impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7; 0, 1, 2, 3, 4, 5, 6, 7; 8);
 impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8; 0, 1, 2, 3, 4, 5, 6, 7, 8; 9);
-impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9;
+impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9; 
                              0, 1, 2, 3, 4, 5, 6, 7, 8, 9; 10);
-impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10;
+impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10; 
                              0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10; 11);
-impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11;
+impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11; 
                              0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11; 12);
-impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12;
+impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12; 
                              0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12; 13);
-impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13;
+impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13; 
                              0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13; 14);
-impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14;
+impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14; 
                              0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14; 15);
 impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15;
                              0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15; 16);
diff --git a/scylla/src/routing.rs b/scylla/src/routing.rs
index 50c7965f..581c83a0 100644
--- a/scylla/src/routing.rs
+++ b/scylla/src/routing.rs
@@ -71,6 +71,26 @@ impl ShardInfo {
             + shard as u16
     }
 
+    /// Returns iterator over source ports `p` such that `shard == shard_of_source_port(p)`.
+    /// Starts at a random port and goes forward by `nr_shards`. After reaching maximum wraps back around.
+    /// Stops once all possibile ports have been returned
+    pub fn iter_source_ports_for_shard(&self, shard: Shard) -> impl Iterator<Item = u16> {
+        assert!(shard < self.nr_shards as u32);
+
+        // Randomly choose a port to start at
+        let starting_port = self.draw_source_port_for_shard(shard);
+
+        // Choose smallest available port number to begin at after wrapping
+        // apply the formula from draw_source_port_for_shard for lowest possible gen_range result
+        let first_valid_port =
+            (49152 + self.nr_shards - 1) / self.nr_shards * self.nr_shards + shard as u16;
+
+        let before_wrap = (starting_port..=65535).step_by(self.nr_shards.into());
+        let after_wrap = (first_valid_port..starting_port).step_by(self.nr_shards.into());
+
+        before_wrap.chain(after_wrap)
+    }
+
     pub fn get_nr_shards(&self) -> u16 {
         self.nr_shards
     }
@@ -108,24 +128,6 @@ impl<'a> TryFrom<&'a HashMap<String, Vec<String>>> for ShardInfo {
     }
 }
 
-#[test]
-fn test_shard_of() {
-    /* Test values taken from the gocql driver.  */
-    let shard_info = ShardInfo::new(0, 4, 12);
-    assert_eq!(
-        shard_info.shard_of(Token {
-            value: -9219783007514621794
-        }),
-        3
-    );
-    assert_eq!(
-        shard_info.shard_of(Token {
-            value: 9222582454147032830
-        }),
-        3
-    );
-}
-
 // An implementation of MurmurHash3 ported from Scylla. Please note that this
 // is not a "correct" implementation of MurmurHash3 - it replicates the same
 // bugs made in the original Cassandra implementation in order to be compatible.
@@ -218,3 +220,63 @@ fn fmix(mut k: Wrapping<i64>) -> Wrapping<i64> {
 
     k
 }
+
+#[cfg(test)]
+mod tests {
+    use super::ShardInfo;
+    use super::Token;
+    use std::collections::HashSet;
+
+    #[test]
+    fn test_shard_of() {
+        /* Test values taken from the gocql driver.  */
+        let shard_info = ShardInfo::new(0, 4, 12);
+        assert_eq!(
+            shard_info.shard_of(Token {
+                value: -9219783007514621794
+            }),
+            3
+        );
+        assert_eq!(
+            shard_info.shard_of(Token {
+                value: 9222582454147032830
+            }),
+            3
+        );
+    }
+
+    #[test]
+    fn test_iter_source_ports_for_shard() {
+        let nr_shards = 4;
+        let max_port_num = 65535;
+        let min_port_num = (49152 + nr_shards - 1) / nr_shards * nr_shards;
+
+        let shard_info = ShardInfo::new(0, nr_shards, 12);
+
+        // Test for each shard
+        for shard in 0..nr_shards {
+            // Find lowest port for this shard
+            let mut lowest_port = min_port_num;
+            while lowest_port % nr_shards != shard {
+                lowest_port += 1;
+            }
+
+            // Find total number of ports the iterator should return
+            let possible_ports_number: usize =
+                ((max_port_num - lowest_port) / nr_shards + 1).into();
+
+            let port_iter = shard_info.iter_source_ports_for_shard(shard.into());
+
+            let mut returned_ports: HashSet<u16> = HashSet::new();
+            for port in port_iter {
+                assert!(!returned_ports.contains(&port)); // No port occurs two times
+                assert!(port % nr_shards == shard); // Each port maps to this shard
+
+                returned_ports.insert(port);
+            }
+
+            // Numbers of ports returned matches the expected value
+            assert_eq!(returned_ports.len(), possible_ports_number);
+        }
+    }
+}
diff --git a/scylla/src/transport/connection.rs b/scylla/src/transport/connection.rs
index 3df204c5..6a23a11c 100644
--- a/scylla/src/transport/connection.rs
+++ b/scylla/src/transport/connection.rs
@@ -39,7 +39,7 @@ pub struct Connection {
     is_shard_aware: bool,
 }
 
-type ResponseHandler = oneshot::Sender<TaskResponse>;
+type ResponseHandler = oneshot::Sender<Result<TaskResponse, QueryError>>;
 
 struct Task {
     request_flags: u8,
@@ -76,12 +76,16 @@ pub struct ConnectionConfig {
     */
 }
 
+// Used to listen for fatal error in connection
+pub type ErrorReceiver = tokio::sync::oneshot::Receiver<QueryError>;
+
 impl Connection {
+    // Returns new connection and ErrorReceiver which can be used to wait for a fatal error
     pub async fn new(
         addr: SocketAddr,
         source_port: Option<u16>,
         config: ConnectionConfig,
-    ) -> Result<Self, std::io::Error> {
+    ) -> Result<(Self, ErrorReceiver), std::io::Error> {
         let stream = match source_port {
             Some(p) => connect_with_source_port(addr, p).await?,
             None => TcpStream::connect(addr).await?,
@@ -92,10 +96,12 @@ impl Connection {
         // TODO: What should be the size of the channel?
         let (sender, receiver) = mpsc::channel(128);
 
-        let (fut, _worker_handle) = Self::router(stream, receiver).remote_handle();
+        let (error_sender, error_receiver) = tokio::sync::oneshot::channel();
+
+        let (fut, _worker_handle) = Self::router(stream, receiver, error_sender).remote_handle();
         tokio::task::spawn(fut);
 
-        Ok(Self {
+        let connection = Connection {
             submit_channel: sender,
             _worker_handle,
             source_port,
@@ -103,7 +109,9 @@ impl Connection {
             shard_info: None,
             config,
             is_shard_aware: false,
-        })
+        };
+
+        Ok((connection, error_receiver))
     }
 
     pub async fn startup(&self, options: HashMap<String, String>) -> Result<Response, QueryError> {
@@ -292,7 +300,7 @@ impl Connection {
                 ErrorKind::Other,
                 "Connection broken",
             )))
-        })?;
+        })??;
         let body_with_ext = frame::parse_response_body_extensions(
             task_response.params.flags,
             self.config.compression,
@@ -310,7 +318,11 @@ impl Connection {
         Ok(response)
     }
 
-    async fn router(mut stream: TcpStream, receiver: mpsc::Receiver<Task>) {
+    async fn router(
+        mut stream: TcpStream,
+        receiver: mpsc::Receiver<Task>,
+        error_sender: tokio::sync::oneshot::Sender<QueryError>,
+    ) {
         let (read_half, write_half) = stream.split();
 
         // Why are using a mutex here?
@@ -330,8 +342,24 @@ impl Connection {
         let r = Self::reader(read_half, &handler_map);
         let w = Self::writer(write_half, &handler_map, receiver);
 
-        // TODO: What to do with this error?
-        let _ = futures::try_join!(r, w);
+        let result = futures::try_join!(r, w);
+
+        let error: QueryError = match result {
+            Ok(_) => return, // Connection was dropped, we can return
+            Err(err) => err,
+        };
+
+        // Respond to all pending requests with the error
+        let response_handlers: HashMap<i16, ResponseHandler> =
+            handler_map.into_inner().unwrap().into_handlers();
+
+        for (_, handler) in response_handlers {
+            // Ignore sending error, request was dropped
+            let _ = handler.send(Err(error.clone()));
+        }
+
+        // If someone is listening for connection errors notify them
+        let _ = error_sender.send(error);
     }
 
     async fn reader(
@@ -366,11 +394,11 @@ impl Connection {
                 // Don't care if sending of the response fails. This must
                 // mean that the receiver side was impatient and is not
                 // waiting for the result anymore.
-                let _ = handler.send(TaskResponse {
+                let _ = handler.send(Ok(TaskResponse {
                     params,
                     opcode,
                     body,
-                });
+                }));
             } else {
                 // Unsolicited frame. This should not happen and indicates
                 // a bug either in the driver, or in the database
@@ -419,7 +447,7 @@ impl Connection {
             .await?;
         }
 
-        Err(std::io::Error::new(ErrorKind::Other, "connection broken").into())
+        Ok(())
     }
 
     pub fn get_shard_info(&self) -> &Option<ShardInfo> {
@@ -453,7 +481,7 @@ pub async fn open_connection(
     addr: SocketAddr,
     source_port: Option<u16>,
     config: ConnectionConfig,
-) -> Result<Connection, QueryError> {
+) -> Result<(Connection, ErrorReceiver), QueryError> {
     open_named_connection(
         addr,
         source_port,
@@ -468,9 +496,10 @@ pub async fn open_named_connection(
     source_port: Option<u16>,
     config: ConnectionConfig,
     driver_name: Option<String>,
-) -> Result<Connection, QueryError> {
+) -> Result<(Connection, ErrorReceiver), QueryError> {
     // TODO: shouldn't all this logic be in Connection::new?
-    let mut connection = Connection::new(addr, source_port, config.clone()).await?;
+    let (mut connection, error_receiver) =
+        Connection::new(addr, source_port, config.clone()).await?;
 
     let options_result = connection.get_options().await?;
 
@@ -522,7 +551,7 @@ pub async fn open_named_connection(
         }
     }
 
-    Ok(connection)
+    Ok((connection, error_receiver))
 }
 
 async fn connect_with_source_port(
@@ -573,6 +602,12 @@ impl ResponseHandlerMap {
         self.stream_set.free(stream_id);
         self.handlers.remove(&stream_id)
     }
+
+    // Retrieves the map of handlers, used after connection breaks
+    // and we have to respond to all of them with an error
+    pub fn into_handlers(self) -> HashMap<i16, ResponseHandler> {
+        self.handlers
+    }
 }
 
 struct StreamIDSet {
diff --git a/scylla/src/transport/connection_keeper.rs b/scylla/src/transport/connection_keeper.rs
index 6455c9b0..0b5b8aac 100644
--- a/scylla/src/transport/connection_keeper.rs
+++ b/scylla/src/transport/connection_keeper.rs
@@ -3,16 +3,18 @@ use crate::routing::ShardInfo;
 use crate::transport::errors::QueryError;
 use crate::transport::{
     connection,
-    connection::{Connection, ConnectionConfig, VerifiedKeyspaceName},
+    connection::{Connection, ConnectionConfig, ErrorReceiver, VerifiedKeyspaceName},
 };
 
 use futures::{future::RemoteHandle, FutureExt};
+use std::io::ErrorKind;
 use std::net::SocketAddr;
 use std::sync::Arc;
 
 /// ConnectionKeeper keeps a Connection to some address and works to keep it open
 pub struct ConnectionKeeper {
     conn_state_receiver: tokio::sync::watch::Receiver<ConnectionState>,
+    use_keyspace_channel: tokio::sync::mpsc::Sender<UseKeyspaceRequest>,
     _worker_handle: RemoteHandle<()>,
 }
 
@@ -32,12 +34,21 @@ struct ConnectionKeeperWorker {
     shard_info_sender: Option<ShardInfoSender>,
     conn_state_sender: tokio::sync::watch::Sender<ConnectionState>,
 
+    // Channel used to receive use keyspace requests
+    use_keyspace_channel: tokio::sync::mpsc::Receiver<UseKeyspaceRequest>,
+
     // Keyspace send in "USE <keyspace name>" when opening each connection
     used_keyspace: Option<VerifiedKeyspaceName>,
 }
 
 pub type ShardInfoSender = Arc<std::sync::Mutex<tokio::sync::watch::Sender<Option<ShardInfo>>>>;
 
+#[derive(Debug)]
+struct UseKeyspaceRequest {
+    keyspace_name: VerifiedKeyspaceName,
+    response_chan: tokio::sync::oneshot::Sender<Result<(), QueryError>>,
+}
+
 impl ConnectionKeeper {
     /// Creates new ConnectionKeeper that starts a connection in the background
     /// # Arguments
@@ -56,12 +67,15 @@ impl ConnectionKeeper {
         let (conn_state_sender, conn_state_receiver) =
             tokio::sync::watch::channel(ConnectionState::Initializing);
 
+        let (use_keyspace_sender, use_keyspace_receiver) = tokio::sync::mpsc::channel(1);
+
         let worker = ConnectionKeeperWorker {
             address,
             config,
             shard_info,
             shard_info_sender,
             conn_state_sender,
+            use_keyspace_channel: use_keyspace_receiver,
             used_keyspace: keyspace_name,
         };
 
@@ -70,6 +84,7 @@ impl ConnectionKeeper {
 
         ConnectionKeeper {
             conn_state_receiver,
+            use_keyspace_channel: use_keyspace_sender,
             _worker_handle: worker_handle,
         }
     }
@@ -113,62 +128,216 @@ impl ConnectionKeeper {
 
     pub async fn use_keyspace(
         &self,
-        keyspace_name: &VerifiedKeyspaceName,
+        keyspace_name: VerifiedKeyspaceName,
     ) -> Result<(), QueryError> {
-        // ConnectionKeeper doesn't have reconnecting yet so this will be ok for now
-        // TODO: Modify once ConnectionKeeper gets reconnecting
+        let (response_sender, response_receiver) = tokio::sync::oneshot::channel();
 
-        self.get_connection()
-            .await?
-            .use_keyspace(keyspace_name)
+        self.use_keyspace_channel
+            .send(UseKeyspaceRequest {
+                keyspace_name,
+                response_chan: response_sender,
+            })
             .await
+            .expect("Bug in ConnectionKeeper::use_keyspace sending");
+        // Other end of this channel is in the Worker, can't be dropped while we have &self to _worker_handle
+
+        response_receiver.await.unwrap() // ConnectionKeeperWorker always responds
     }
 }
 
+enum RunConnectionRes {
+    // An error occured during the connection
+    Error(QueryError),
+    // ConnectionKeeper was dropped and channels closed, we should stop
+    ShouldStop,
+}
+
 impl ConnectionKeeperWorker {
-    pub async fn work(self) {
-        let cur_connection = self.open_new_connection().await;
-
-        match &cur_connection {
-            Ok(conn) => {
-                let _ = self
-                    .conn_state_sender
-                    .send(ConnectionState::Connected(conn.clone()));
-
-                let new_shard_info: Option<ShardInfo> = conn.get_shard_info().clone();
-
-                if let Some(sender) = &self.shard_info_sender {
-                    // Ignore sending error
-                    // If no one wants to get shard_info that's OK
-                    // If lock is poisoned do nothing
-                    if let Ok(sender_locked) = sender.lock() {
-                        let _ = sender_locked.send(new_shard_info);
-                    }
-                }
+    pub async fn work(mut self) {
+        // Reconnect at most every 8 seconds
+        let reconnect_cooldown = tokio::time::Duration::from_secs(8);
+        let mut last_reconnect_time;
+
+        loop {
+            last_reconnect_time = tokio::time::Instant::now();
+
+            // Connect and wait for error
+            let current_error: QueryError = match self.run_connection().await {
+                RunConnectionRes::Error(e) => e,
+                RunConnectionRes::ShouldStop => return,
+            };
+
+            // Mark the connection as broken, wait cooldown and reconnect
+            if self
+                .conn_state_sender
+                .send(ConnectionState::Broken(current_error))
+                .is_err()
+            {
+                // ConnectionKeeper was dropped, we should stop
+                return;
             }
-            Err(e) => {
-                let _ = self
-                    .conn_state_sender
-                    .send(ConnectionState::Broken(e.clone()));
-            } // TODO: Wait for connection to fail, then create new, loop it
-        };
+
+            let next_reconnect_time = last_reconnect_time
+                .checked_add(reconnect_cooldown)
+                .unwrap_or_else(tokio::time::Instant::now);
+
+            tokio::time::sleep_until(next_reconnect_time).await;
+        }
     }
 
-    async fn open_new_connection(&self) -> Result<Arc<Connection>, QueryError> {
-        let mut source_port: Option<u16> = None;
-        if let Some(info) = &self.shard_info {
-            source_port = Some(info.draw_source_port_for_shard(info.shard.into()));
+    // Opens a new connection and waits until some fatal error occurs
+    async fn run_connection(&mut self) -> RunConnectionRes {
+        // Connect to the node
+        let (connection, mut error_receiver) = match self.open_new_connection().await {
+            Ok(opened) => opened,
+            Err(e) => return RunConnectionRes::Error(e),
+        };
+
+        // Mark connection as Connected
+        if self
+            .conn_state_sender
+            .send(ConnectionState::Connected(connection.clone()))
+            .is_err()
+        {
+            // If the channel was dropped we should stop
+            return RunConnectionRes::ShouldStop;
         }
 
-        let new_conn =
-            connection::open_connection(self.address, source_port, self.config.clone()).await?;
+        // Notify about new shard info
+        if let Some(sender) = &self.shard_info_sender {
+            let new_shard_info: Option<ShardInfo> = connection.get_shard_info().clone();
+
+            // Ignore sending error
+            // If no one wants to get shard_info that's OK
+            // If lock is poisoned do nothing
+            if let Ok(sender_locked) = sender.lock() {
+                let _ = sender_locked.send(new_shard_info);
+            }
+        }
 
+        // Use the specified keyspace
         if let Some(keyspace_name) = &self.used_keyspace {
-            let _ = new_conn.use_keyspace(&keyspace_name).await;
+            let _ = connection.use_keyspace(&keyspace_name).await;
             // Ignore the error, used_keyspace could be set a long time ago and then deleted
             // user gets all errors from session.use_keyspace()
         }
 
-        Ok(Arc::new(new_conn))
+        let connection_closed_error = QueryError::IOError(Arc::new(std::io::Error::new(
+            ErrorKind::Other,
+            "Connection closed",
+        )));
+
+        // Wait for events - a use keyspace request or a fatal error
+        loop {
+            tokio::select! {
+                recv_res = self.use_keyspace_channel.recv() => {
+                    match recv_res {
+                        Some(request) => {
+                            self.used_keyspace = Some(request.keyspace_name.clone());
+
+                            // Send USE KEYSPACE request, send result if channel wasn't closed
+                            let res = connection.use_keyspace(&request.keyspace_name).await;
+                            let _ = request.response_chan.send(res);
+                        },
+                        None => return RunConnectionRes::ShouldStop, // If the channel was dropped we should stop
+                    }
+                },
+                connection_error = &mut error_receiver => {
+                    let error = connection_error.unwrap_or(connection_closed_error);
+                    return RunConnectionRes::Error(error);
+                }
+            }
+        }
+    }
+
+    async fn open_new_connection(&self) -> Result<(Arc<Connection>, ErrorReceiver), QueryError> {
+        let (connection, error_receiver) = match &self.shard_info {
+            Some(info) => self.open_new_connection_to_shard(info).await?,
+            None => connection::open_connection(self.address, None, self.config.clone()).await?,
+        };
+
+        Ok((Arc::new(connection), error_receiver))
+    }
+
+    async fn open_new_connection_to_shard(
+        &self,
+        shard_info: &ShardInfo,
+    ) -> Result<(Connection, ErrorReceiver), QueryError> {
+        // Create iterator over all possible source ports for this shard
+        let source_port_iter = shard_info.iter_source_ports_for_shard(shard_info.shard.into());
+
+        for port in source_port_iter {
+            let connect_result =
+                connection::open_connection(self.address, Some(port), self.config.clone()).await;
+
+            match connect_result {
+                Err(err) if err.is_address_in_use() => continue, // If port collision happened try next port
+                result => return result,
+            }
+        }
+
+        // Tried all source ports for that shard, give up
+        return Err(QueryError::IOError(Arc::new(std::io::Error::new(
+            std::io::ErrorKind::AddrInUse,
+            "Could not find free source port for shard",
+        ))));
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::ConnectionKeeper;
+    use crate::transport::connection::ConnectionConfig;
+    use std::net::{SocketAddr, ToSocketAddrs};
+
+    // Open many connections to a node
+    // Port collision should occur
+    // If they are not handled this test will most likely fail
+    #[tokio::test]
+    #[ignore]
+    async fn many_connections() {
+        let connections_number = 512;
+
+        let connect_address: SocketAddr = std::env::var("SCYLLA_URI")
+            .unwrap_or_else(|_| "127.0.0.1:9042".to_string())
+            .to_socket_addrs()
+            .unwrap()
+            .next()
+            .unwrap();
+
+        let connection_config = ConnectionConfig {
+            compression: None,
+            tcp_nodelay: false,
+        };
+
+        // Get shard info from a single connection, all connections will open to this shard
+        let conn_keeper =
+            ConnectionKeeper::new(connect_address, connection_config.clone(), None, None, None);
+        let shard_info = conn_keeper
+            .get_connection()
+            .await
+            .unwrap()
+            .get_shard_info()
+            .clone();
+
+        // Open the connections
+        let mut conn_keepers: Vec<ConnectionKeeper> = Vec::new();
+
+        for _ in 0..connections_number {
+            let conn_keeper = ConnectionKeeper::new(
+                connect_address,
+                connection_config.clone(),
+                shard_info.clone(),
+                None,
+                None,
+            );
+
+            conn_keepers.push(conn_keeper);
+        }
+
+        // Check that each connection keeper connected succesfully
+        for conn_keeper in conn_keepers {
+            conn_keeper.get_connection().await.unwrap();
+        }
     }
 }
diff --git a/scylla/src/transport/errors.rs b/scylla/src/transport/errors.rs
index b4c18d37..1d20f999 100644
--- a/scylla/src/transport/errors.rs
+++ b/scylla/src/transport/errors.rs
@@ -1,5 +1,6 @@
 use crate::frame::frame_errors::{FrameError, ParseError};
 use crate::frame::value::SerializeValuesError;
+use std::io::ErrorKind;
 use std::sync::Arc;
 use thiserror::Error;
 
@@ -146,3 +147,18 @@ impl From<BadKeyspaceName> for QueryError {
         QueryError::BadQuery(BadQuery::BadKeyspaceName(keyspace_err))
     }
 }
+
+impl QueryError {
+    /// Checks if this query error is an IO Error of kind AddrInUse
+    /// This type of error often occurs when we are choosing random port numbers
+    /// and there is a collision.
+    pub(crate) fn is_address_in_use(&self) -> bool {
+        if let QueryError::IOError(io_error) = self {
+            if io_error.kind() == ErrorKind::AddrInUse {
+                return true;
+            }
+        }
+
+        false
+    }
+}
diff --git a/scylla/src/transport/node.rs b/scylla/src/transport/node.rs
index 3864125a..718324f1 100644
--- a/scylla/src/transport/node.rs
+++ b/scylla/src/transport/node.rs
@@ -274,12 +274,12 @@ impl NodeWorker {
 
         match &*node_conns {
             NodeConnections::Single(conn_keeper) => {
-                let fut = conn_keeper.use_keyspace(keyspace_name);
+                let fut = conn_keeper.use_keyspace(keyspace_name.clone());
                 use_keyspace_futures.push(fut);
             }
             NodeConnections::Sharded { shard_conns, .. } => {
                 for conn_keeper in shard_conns {
-                    let fut = conn_keeper.use_keyspace(keyspace_name);
+                    let fut = conn_keeper.use_keyspace(keyspace_name.clone());
                     use_keyspace_futures.push(fut);
                 }
             }


Test patch:
diff --git a/scylla/src/frame/cql_collections_test.rs b/scylla/src/frame/cql_collections_test.rs
index 589f9150..87f6c298 100644
--- a/scylla/src/frame/cql_collections_test.rs
+++ b/scylla/src/frame/cql_collections_test.rs
@@ -3,7 +3,9 @@ use crate::{IntoTypedRows, SessionBuilder};
 use std::collections::HashMap;
 use std::env;
 
+// TODO: Requires a running local Scylla instance
 #[tokio::test]
+#[ignore]
 async fn test_cql_collections() {
     // Create connection
     let uri = env::var("SCYLLA_URI").unwrap_or_else(|_| "127.0.0.1:9042".to_string());
diff --git a/scylla/src/frame/value_tests.rs b/scylla/src/frame/value_tests.rs
index 76d3671e..0d923164 100644
--- a/scylla/src/frame/value_tests.rs
+++ b/scylla/src/frame/value_tests.rs
@@ -1,6 +1,6 @@
 use super::value::{
-    BatchValues, MaybeUnset, SerializeValuesError, SerializedResult, SerializedValues, Unset,
-    Value, ValueList, ValueTooBig,
+    BatchValues, SerializeValuesError, SerializedResult, SerializedValues, Value, ValueList,
+    ValueTooBig,
 };
 use bytes::BufMut;
 use std::borrow::Cow;
@@ -33,17 +33,6 @@ fn option_value() {
     assert_eq!(serialized(null_i32), &(-1_i32).to_be_bytes()[..]);
 }
 
-#[test]
-fn unset_value() {
-    assert_eq!(serialized(Unset), &(-2_i32).to_be_bytes()[..]);
-
-    let unset_i32: MaybeUnset<i32> = MaybeUnset::Unset;
-    assert_eq!(serialized(unset_i32), &(-2_i32).to_be_bytes()[..]);
-
-    let set_i32: MaybeUnset<i32> = MaybeUnset::Set(32);
-    assert_eq!(serialized(set_i32), vec![0, 0, 0, 4, 0, 0, 0, 32]);
-}
-
 #[test]
 fn ref_value() {
     assert_eq!(serialized(&1_i32), serialized(1_i32));
diff --git a/scylla/src/transport/session_test.rs b/scylla/src/transport/session_test.rs
index 1d9bc95c..3f198873 100644
--- a/scylla/src/transport/session_test.rs
+++ b/scylla/src/transport/session_test.rs
@@ -3,7 +3,9 @@ use crate::routing::hash3_x64_128;
 use crate::transport::errors::{BadKeyspaceName, BadQuery, QueryError};
 use crate::{IntoTypedRows, Session, SessionBuilder};
 
+// TODO: Requires a running local Scylla instance
 #[tokio::test]
+#[ignore]
 async fn test_unprepared_statement() {
     let uri = std::env::var("SCYLLA_URI").unwrap_or_else(|_| "127.0.0.1:9042".to_string());
     let session = SessionBuilder::new().known_node(uri).build().await.unwrap();
@@ -62,6 +64,7 @@ async fn test_unprepared_statement() {
 }
 
 #[tokio::test]
+#[ignore]
 async fn test_prepared_statement() {
     let uri = std::env::var("SCYLLA_URI").unwrap_or_else(|_| "127.0.0.1:9042".to_string());
     let session = SessionBuilder::new().known_node(uri).build().await.unwrap();
@@ -177,6 +180,7 @@ async fn test_prepared_statement() {
 }
 
 #[tokio::test]
+#[ignore]
 async fn test_batch() {
     let uri = std::env::var("SCYLLA_URI").unwrap_or_else(|_| "127.0.0.1:9042".to_string());
     let session = SessionBuilder::new().known_node(uri).build().await.unwrap();
@@ -241,6 +245,7 @@ async fn test_batch() {
 }
 
 #[tokio::test]
+#[ignore]
 async fn test_token_calculation() {
     let uri = std::env::var("SCYLLA_URI").unwrap_or_else(|_| "127.0.0.1:9042".to_string());
     let session = SessionBuilder::new().known_node(uri).build().await.unwrap();