<commit_before>// Copyright 2017 click2stream, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::{HashMap, VecDeque};
use std::net::SocketAddr;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::time::Duration;

use bytes::BytesMut;

use futures::task::{Context, Poll, Waker};
use futures::{Future, Stream};

use tokio::io::{AsyncRead, AsyncWrite};
use tokio::net::TcpStream;

use crate::context::ApplicationContext;
use crate::net::arrow::error::{ArrowError, ConnectionError};
use crate::net::arrow::proto::msg::control::{
    ControlMessageFactory, EC_CONNECTION_ERROR, EC_NO_ERROR,
};
use crate::net::arrow::proto::msg::ArrowMessage;
use crate::svc_table::{BoxServiceTable, ServiceTable};
use crate::utils::logger::{BoxLogger, Logger};

const OUTPUT_BUFFER_LIMIT: usize = 4 * 1024 * 1024;

const CONNECTION_TIMEOUT: Duration = Duration::from_secs(20);

/// Session context.
struct SessionContext {
    service_id: u16,
    session_id: u32,
    input: BytesMut,
    output: BytesMut,
    session_manager_task: Option<Waker>,
    session_transport_task: Option<Waker>,
    closed: bool,
    error: Option<ConnectionError>,
}

impl SessionContext {
    /// Create a new session context for a given service ID and session ID.
    fn new(service_id: u16, session_id: u32) -> Self {
        Self {
            service_id,
            session_id,
            input: BytesMut::with_capacity(8192),
            output: BytesMut::with_capacity(8192),
            session_manager_task: None,
            session_transport_task: None,
            closed: false,
            error: None,
        }
    }

    /// Extend the output buffer with data from a given Arrow Message.
    fn push_output_message(&mut self, msg: ArrowMessage) {
        // ignore all incoming messages after the connection gets closed
        if self.closed {
            return;
        }

        let data = msg.payload();

        if (self.output.len() + data.len()) > OUTPUT_BUFFER_LIMIT {
            // we cannot backpressure here, so we'll set an error state
            self.set_error(ConnectionError::new("output buffer limit exceeded"));
        } else {
            self.output.extend_from_slice(data);

            // we MUST notify the transport task that there is some data
            // available in the output buffer again
            if !self.output.is_empty() {
                if let Some(task) = self.session_transport_task.take() {
                    task.wake();
                }
            }
        }
    }

    /// Take all the data from the input buffer and return them as an Arrow
    /// Message. The method returns:
    /// * `Poll::Ready(Some(Ok(_)))` if there was some data available
    /// * `Poll::Ready(None)` if there was no data available and the context
    ///   has been closed
    /// * `Poll::Pending` if there was no data available
    fn take_input_message(
        &mut self,
        cx: &mut Context,
    ) -> Poll<Option<Result<ArrowMessage, ConnectionError>>> {
        let data = self.input.split().freeze();

        // we MUST notify the transport task that the input buffer is empty
        // again
        if let Some(task) = self.session_transport_task.take() {
            task.wake();
        }

        if !data.is_empty() {
            let message = ArrowMessage::new(self.service_id, self.session_id, data);

            Poll::Ready(Some(Ok(message)))
        } else if self.closed {
            match self.error.take() {
                Some(err) => Poll::Ready(Some(Err(err))),
                None => Poll::Ready(None),
            }
        } else {
            // save the current task and wait until there is some data
            // available in the input buffer again
            self.session_manager_task = Some(cx.waker().clone());

            Poll::Pending
        }
    }

    /// Try to read data from a given `AsyncRead` stream into the input buffer.
    fn poll_read_input<S>(&mut self, cx: &mut Context, stream: &mut S) -> Poll<()>
    where
        S: AsyncRead + Unpin,
    {
        let input_buffer_len = self.input.len();
        let input_buffer_capacity = self.input.capacity();

        let stream = Pin::new(stream);

        if self.closed {
            Poll::Ready(())
        } else if input_buffer_len >= input_buffer_capacity {
            // save the current task and wait until there is some space in
            // the input buffer again
            self.session_transport_task = Some(cx.waker().clone());

            Poll::Pending
        } else if let Poll::Ready(res) = stream.poll_read_buf(cx, &mut self.input) {
            match res {
                Ok(len) if len == 0 => self.close(),
                Err(err) => self.set_error(ConnectionError::from(err)),
                _ => (),
            }
<commit_msg>Fix input buffer capacity check within the session connection<commit_after>// Copyright 2017 click2stream, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::{HashMap, VecDeque};
use std::net::SocketAddr;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::time::Duration;

use bytes::BytesMut;

use futures::task::{Context, Poll, Waker};
use futures::{Future, Stream};

use tokio::io::{AsyncRead, AsyncWrite};
use tokio::net::TcpStream;

use crate::context::ApplicationContext;
use crate::net::arrow::error::{ArrowError, ConnectionError};
use crate::net::arrow::proto::msg::control::{
    ControlMessageFactory, EC_CONNECTION_ERROR, EC_NO_ERROR,
};
use crate::net::arrow::proto::msg::ArrowMessage;
use crate::svc_table::{BoxServiceTable, ServiceTable};
use crate::utils::logger::{BoxLogger, Logger};

const OUTPUT_BUFFER_LIMIT: usize = 4 * 1024 * 1024;

const CONNECTION_TIMEOUT: Duration = Duration::from_secs(20);

/// Session context.
struct SessionContext {
    service_id: u16,
    session_id: u32,
    input: BytesMut,
    output: BytesMut,
    buffer_capacity: usize,
    session_manager_task: Option<Waker>,
    session_transport_task: Option<Waker>,
    closed: bool,
    error: Option<ConnectionError>,
}

impl SessionContext {
    /// Create a new session context for a given service ID and session ID.
    fn new(service_id: u16, session_id: u32) -> Self {
        let buffer_capacity = 8192;

        Self {
            service_id,
            session_id,
            input: BytesMut::with_capacity(buffer_capacity),
            output: BytesMut::with_capacity(buffer_capacity),
            buffer_capacity,
            session_manager_task: None,
            session_transport_task: None,
            closed: false,
            error: None,
        }
    }

    /// Extend the output buffer with data from a given Arrow Message.
    fn push_output_message(&mut self, msg: ArrowMessage) {
        // ignore all incoming messages after the connection gets closed
        if self.closed {
            return;
        }

        let data = msg.payload();

        if (self.output.len() + data.len()) > OUTPUT_BUFFER_LIMIT {
            // we cannot backpressure here, so we'll set an error state
            self.set_error(ConnectionError::new("output buffer limit exceeded"));
        } else {
            self.output.extend_from_slice(data);

            // we MUST notify the transport task that there is some data
            // available in the output buffer again
            if !self.output.is_empty() {
                if let Some(task) = self.session_transport_task.take() {
                    task.wake();
                }
            }
        }
    }

    /// Take all the data from the input buffer and return them as an Arrow
    /// Message. The method returns:
    /// * `Poll::Ready(Some(Ok(_)))` if there was some data available
    /// * `Poll::Ready(None)` if there was no data available and the context
    ///   has been closed
    /// * `Poll::Pending` if there was no data available
    fn take_input_message(
        &mut self,
        cx: &mut Context,
    ) -> Poll<Option<Result<ArrowMessage, ConnectionError>>> {
        let data = self.input.split().freeze();

        // we MUST notify the transport task that the input buffer is empty
        // again
        if let Some(task) = self.session_transport_task.take() {
            task.wake();
        }

        if !data.is_empty() {
            let message = ArrowMessage::new(self.service_id, self.session_id, data);

            Poll::Ready(Some(Ok(message)))
        } else if self.closed {
            match self.error.take() {
                Some(err) => Poll::Ready(Some(Err(err))),
                None => Poll::Ready(None),
            }
        } else {
            // save the current task and wait until there is some data
            // available in the input buffer again
            self.session_manager_task = Some(cx.waker().clone());

            Poll::Pending
        }
    }

    /// Try to read data from a given `AsyncRead` stream into the input buffer.
    fn poll_read_input<S>(&mut self, cx: &mut Context, stream: &mut S) -> Poll<()>
    where
        S: AsyncRead + Unpin,
    {
        let input_buffer_len = self.input.len();

        let stream = Pin::new(stream);

        if self.closed {
            Poll::Ready(())
        } else if input_buffer_len >= self.buffer_capacity {
            // save the current task and wait until there is some space in
            // the input buffer again
            self.session_transport_task = Some(cx.waker().clone());

            Poll::Pending
        } else if let Poll::Ready(res) = stream.poll_read_buf(cx, &mut self.input) {
            match res {
                Ok(len) if len == 0 => self.close(),
                Err(err) => self.set_error(ConnectionError::from(err)),
                _ => (),
            }

{'fraction_non_alphanumeric': 0.1142750790991997, 'fraction_numerical': 0.0065140517401823935, 'mean_word_length': 2.383816120906801, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 39, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5491342', 'n_tokens_mistral': 3034, 'n_tokens_neox': 2964, 'n_words': 931}