<commit_before>use super::{read_query, wrap_query, write_query};
use crate::{
    Document, ReqlResponse, Request, Result,
    Session, SessionManager,
    errors::*,
};

use protobuf::ProtobufEnum;
use ql2::proto::{Query_QueryType as QueryType, Response_ErrorType as ErrorType,
                 Response_ResponseType as ResponseType};
use r2d2::PooledConnection;
use serde::de::DeserializeOwned;
use serde_json::{Value, from_slice, from_value};
use std::error::Error as StdError;
use futures::{Future, Sink};

impl<T: DeserializeOwned + Send + std::fmt::Debug> Request<T> {
    fn conn(&self) -> Result<PooledConnection<SessionManager>> {
        match self.pool.get() {
            Ok(mut conn) => {
                conn.id = conn.id.wrapping_add(1);
                Ok(conn)
            }
            Err(error) => Err(error)?,
        }
    }

    pub fn submit(mut self) {
        let mut conn = match self.conn() {
            Ok(conn) => conn,
            Err(error) => {
                let _ = self.tx.clone().send(Err(error));
                return;
            }
        };
        // Try sending the query
        log::debug!("submiting to server");
        {
            let mut i = 0;
            let mut connect = false;
            let reproducible = self.cfg.opts.reproducible;
            while i < self.cfg.opts.retries || reproducible {
                log::debug!("attempt number {}", i+1);
                // Open a new connection if necessary
                if connect {
                    log::debug!("reconnecting...");
                    drop(&mut conn);
                    conn = match self.conn() {
                        Ok(c) => c,
                        Err(error) => {
                            if i == self.cfg.opts.retries - 1 {
                                let _ = self.tx.clone().send(Err(error.into()));
                                if !reproducible {
                                    return;
                                }
                            }
                            i += 1;
                            continue;
                        }
                    };
                    self.write = true;
                }
                let commands = {
                    let term = self.term.clone();
                    self.encode(&term, false)
                };
                let opts = {
                    let term = self.opts.clone();
                    let res = self.encode(&term, true);
                    if res.is_empty() {
                        None
                    } else {
                        log::debug!("{}", res);
                        Some(res)
                    }
                };
                let mut query = wrap_query(QueryType::START, Some(commands), opts);
                log::debug!("{}", query);
                // Submit the query if necessary
                if self.write || reproducible {
                    log::debug!("submitting query");
                    if let Err(error) = write_query(&mut conn, &query) {
                        connect = true;
                        if i == self.cfg.opts.retries - 1 {
                            let _ = self.tx.clone().send(Err(error.into()));
                            if !reproducible {
                                return;
                            }
                        }
                        i += 1;
                        continue;
                    }
                    if reproducible {
                        connect = true;
                    } else {
                        connect = false;
                    }
                }
                // Handle the response
                let mut more = true;
                while more {
                    match self.process(&mut conn) {
                        Err(error) => {
                            if i == self.cfg.opts.retries - 1 || !self.retry {
                                let _ = self.tx.clone().send(Err(error.into()));
                                if !reproducible {
                                    return;
                                }
                            }
                            i += 1;
                            //continue;
                            more = false;
                            break;
                        },
                        Ok(true) => { // Continue processing, for example in changefeeds
                            query = wrap_query(QueryType::CONTINUE, None, None);
                            if let Err(error) = write_query(&mut conn, &mut query) {
                                self.write = true;
                                self.retry = true;
                                let _ = self.tx.clone().send(Err(error.into()));<commit_msg>Forgot some wait() from when regular channels were tested<commit_after>use super::{read_query, wrap_query, write_query};
use crate::{
    Document, ReqlResponse, Request, Result,
    Session, SessionManager,
    errors::*,
};

use protobuf::ProtobufEnum;
use ql2::proto::{Query_QueryType as QueryType, Response_ErrorType as ErrorType,
                 Response_ResponseType as ResponseType};
use r2d2::PooledConnection;
use serde::de::DeserializeOwned;
use serde_json::{Value, from_slice, from_value};
use std::error::Error as StdError;
use futures::{Future, Sink};

impl<T: DeserializeOwned + Send + std::fmt::Debug> Request<T> {
    fn conn(&self) -> Result<PooledConnection<SessionManager>> {
        match self.pool.get() {
            Ok(mut conn) => {
                conn.id = conn.id.wrapping_add(1);
                Ok(conn)
            }
            Err(error) => Err(error)?,
        }
    }

    pub fn submit(mut self) {
        let mut conn = match self.conn() {
            Ok(conn) => conn,
            Err(error) => {
                let _ = self.tx.clone().send(Err(error)).wait();
                return;
            }
        };
        // Try sending the query
        log::debug!("submiting to server");
        {
            let mut i = 0;
            let mut connect = false;
            let reproducible = self.cfg.opts.reproducible;
            while i < self.cfg.opts.retries || reproducible {
                log::debug!("attempt number {}", i+1);
                // Open a new connection if necessary
                if connect {
                    log::debug!("reconnecting...");
                    drop(&mut conn);
                    conn = match self.conn() {
                        Ok(c) => c,
                        Err(error) => {
                            if i == self.cfg.opts.retries - 1 {
                                let _ = self.tx.clone().send(Err(error.into())).wait();
                                if !reproducible {
                                    return;
                                }
                            }
                            i += 1;
                            continue;
                        }
                    };
                    self.write = true;
                }
                let commands = {
                    let term = self.term.clone();
                    self.encode(&term, false)
                };
                let opts = {
                    let term = self.opts.clone();
                    let res = self.encode(&term, true);
                    if res.is_empty() {
                        None
                    } else {
                        log::debug!("{}", res);
                        Some(res)
                    }
                };
                let mut query = wrap_query(QueryType::START, Some(commands), opts);
                log::debug!("{}", query);
                // Submit the query if necessary
                if self.write || reproducible {
                    log::debug!("submitting query");
                    if let Err(error) = write_query(&mut conn, &query) {
                        connect = true;
                        if i == self.cfg.opts.retries - 1 {
                            let _ = self.tx.clone().send(Err(error.into())).wait();
                            if !reproducible {
                                return;
                            }
                        }
                        i += 1;
                        continue;
                    }
                    if reproducible {
                        connect = true;
                    } else {
                        connect = false;
                    }
                }
                // Handle the response
                let mut more = true;
                while more {
                    match self.process(&mut conn) {
                        Err(error) => {
                            if i == self.cfg.opts.retries - 1 || !self.retry {
                                let _ = self.tx.clone().send(Err(error.into())).wait();
                                if !reproducible {
                                    return;
                                }
                            }
                            i += 1;
                            //continue;
                            more = false;
                            break;
                        },
                        Ok(true) => { // Continue processing, for example in changefeeds
                            query = wrap_query(QueryType::CONTINUE, None, None);
                            if let Err(error) = write_query(&mut conn, &mut query) {
                                self.write = true;
                                self.retry = true;
                                let _ = self.tx.clone().send(Err(error.into())).wait();
{'fraction_non_alphanumeric': 0.10454875142472282, 'fraction_numerical': 0.0024867889337892445, 'mean_word_length': 0.8836846213895394, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 29, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 18, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5276401', 'n_tokens_mistral': 2352, 'n_tokens_neox': 2190, 'n_words': 505}