<commit_before>//! Test data production and consumption using high
//! level clients (StreamConsumer and FutureProducer).
extern crate env_logger;
extern crate futures;
extern crate rand;
extern crate rdkafka;

use futures::*;

use rdkafka::{Message, Timestamp};
use rdkafka::consumer::{Consumer, CommitMode};
use rdkafka::topic_partition_list::{Offset, TopicPartitionList};
use rdkafka::error::KafkaError;

mod utils;
use utils::*;

use std::time::{Duration, Instant};
use std::collections::HashMap;

// All messages should go to the same partition.
#[test]
fn test_produce_partition() {
    let _r = env_logger::init();

    let topic_name = rand_test_topic();
    let message_map = produce_messages(&topic_name, 100, &value_fn, &key_fn, Some(0), None);

    let res = message_map.iter()
        .filter(|&(&(partition, _), _)| partition == 0)
        .count();

    assert_eq!(res, 100);
}

// All produced messages should be consumed.
#[test]
fn test_produce_consume_base() {
    let _r = env_logger::init();

    let topic_name = rand_test_topic();
    let message_map = produce_messages(&topic_name, 100, &value_fn, &key_fn, None, None);
    let consumer = create_stream_consumer(&rand_test_group(), None);
    consumer.subscribe(&[topic_name.as_str()]).unwrap();

    let _consumer_future = consumer.start()
        .take(100)
        .for_each(|message| {
            match message {
                Ok(m) => {
                    let id = message_map[&(m.partition(), m.offset())];
                    match m.timestamp() {
                        Timestamp::CreateTime(timestamp) => assert!(timestamp > 1489495183000),
                        _ => panic!("Expected createtime for message timestamp")
                    };
                    assert_eq!(m.payload_view::<str>().unwrap().unwrap(), value_fn(id));
                    assert_eq!(m.key_view::<str>().unwrap().unwrap(), key_fn(id));
                    assert_eq!(m.topic(), topic_name.as_str());
                },
                Err(e) => panic!("Error receiving message: {:?}", e)
            };
            Ok(())
        })
        .wait();
}

// All produced messages should be consumed.
#[test]
fn test_produce_consume_base_assign() {
    let _r = env_logger::init();

    let topic_name = rand_test_topic();
    produce_messages(&topic_name, 10, &value_fn, &key_fn, Some(0), None);
    produce_messages(&topic_name, 10, &value_fn, &key_fn, Some(1), None);
    produce_messages(&topic_name, 10, &value_fn, &key_fn, Some(2), None);
    let consumer = create_stream_consumer(&rand_test_group(), None);
    let mut tpl = TopicPartitionList::new();
    tpl.add_partition_offset(&topic_name, 0, Offset::Beginning);
    tpl.add_partition_offset(&topic_name, 1, Offset::Offset(2));
    tpl.add_partition_offset(&topic_name, 2, Offset::Offset(9));
    consumer.assign(&tpl).unwrap();

    let mut partition_count = vec![0, 0, 0];

    let _consumer_future = consumer.start()
        .take(19)
        .for_each(|message| {
            match message {
                Ok(m) => partition_count[m.partition() as usize] += 1,
                Err(e) => panic!("Error receiving message: {:?}", e)
            };
            Ok(())
        })
        .wait();

    assert_eq!(partition_count, vec![10, 8, 1]);
}

// All produced messages should be consumed.
#[test]
fn test_produce_consume_with_timestamp() {
    let _r = env_logger::init();

    let topic_name = rand_test_topic();
    let message_map = produce_messages(&topic_name, 100, &value_fn, &key_fn, Some(0), Some(1111));
    let consumer = create_stream_consumer(&rand_test_group(), None);
    consumer.subscribe(&[topic_name.as_str()]).unwrap();

    let _consumer_future = consumer.start()
        .take(100)
        .for_each(|message| {
            match message {
                Ok(m) => {
                    let id = message_map[&(m.partition(), m.offset())];
                    assert_eq!(m.timestamp(), Timestamp::CreateTime(1111));
                    assert_eq!(m.payload_view::<str>().unwrap().unwrap(), value_fn(id));
                    assert_eq!(m.key_view::<str>().unwrap().unwrap(), key_fn(id));
                },
                Err(e) => panic!("Error receiving message: {:?}", e)
            };
            Ok(())
        })
        .wait();

    produce_messages(&topic_name, 10, &value_fn, &key_fn, Some(0), Some(999999));

    // Lookup the offsets
    let tpl = consumer.offsets_for_timestamp(999999, 10000).unwrap();
    let tp = tpl.find_partition(&topic_name, 0).unwrap();
    assert_eq!(tp.topic(), topic_name);
    assert_eq!(tp.offset(), Offset::Offset(100));
    assert_eq!(tp.partition(), 0);
    assert_eq!(tp.error(), Ok(()));
}

#[test]
fn test_consume_with_no_message_error() {
    let _r = env_logger::init();

    let consumer = create_stream_consumer(&rand_test_group(), None);

    let message_stream = consumer.start_with(Duration::from_millis(200), true);

    let mut first_poll_time = None;
    let mut timeouts_count = 0;
    for message in message_stream.wait() {
        match message {
            Ok(Err(KafkaError::NoMessageReceived)) => {
                // TODO: use entry interface for Options once available
                if first_poll_time.is_none() {
                    first_poll_time = Some(Instant::now());
                }
                timeouts_count += 1;
                if timeouts_count == 26 {
                    break;
                }
            }
            Ok(m) => panic!("A message was actually received: {:?}", m),
            Err(e) => panic!("Unexpected error while receiving message: {:?}", e)
        };
    }

    assert_eq!(timeouts_count, 26);
    // It should take 5000ms
    println!("Duration: {:?}", first_poll_time.unwrap().elapsed());
    assert!(first_poll_time.unwrap().elapsed() < Duration::from_millis(7000));
    assert!(first_poll_time.unwrap().elapsed() > Duration::from_millis(4500));
}

// TODO: add check that commit cb gets called correctly
#[test]
fn test_consumer_commit_message() {
    let _r = env_logger::init();

    let topic_name = rand_test_topic();
    produce_messages(&topic_name, 10, &value_fn, &key_fn, Some(0), None);
    produce_messages(&topic_name, 11, &value_fn, &key_fn, Some(1), None);
    produce_messages(&topic_name, 12, &value_fn, &key_fn, Some(2), None);
    let consumer = create_stream_consumer(&rand_test_group(), None);
    consumer.subscribe(&[topic_name.as_str()]).unwrap();

    let _consumer_future = consumer.start()
        .take(33)
        .for_each(|message| {
            match message {
                Ok(m) => {
                    if m.partition() == 1 {
                        consumer.commit_message(&m, CommitMode::Async).unwrap();
                    }
                },
                Err(e) => panic!("error receiving message: {:?}", e)
            };
            Ok(())
        })
        .wait();

    assert_eq!(consumer.fetch_watermarks(&topic_name, 0, 5000).unwrap(), (0, 10));
    assert_eq!(consumer.fetch_watermarks(&topic_name, 1, 5000).unwrap(), (0, 11));
    assert_eq!(consumer.fetch_watermarks(&topic_name, 2, 5000).unwrap(), (0, 12));

    let mut assignment = TopicPartitionList::new();
    assignment.add_partition_offset(&topic_name, 0, Offset::Invalid);
    assignment.add_partition_offset(&topic_name, 1, Offset::Invalid);
    assignment.add_partition_offset(&topic_name, 2, Offset::Invalid);
    assert_eq!(assignment, consumer.assignment().unwrap());

    let mut committed = TopicPartitionList::new();
    committed.add_partition_offset(&topic_name, 0, Offset::Invalid);
    committed.add_partition_offset(&topic_name, 1, Offset::Offset(11));
    committed.add_partition_offset(&topic_name, 2, Offset::Invalid);
    assert_eq!(committed, consumer.committed(5000).unwrap());

    let mut position = TopicPartitionList::new();
    position.add_partition_offset(&topic_name, 0, Offset::Offset(10));
    position.add_partition_offset(&topic_name, 1, Offset::Offset(11));
    position.add_partition_offset(&topic_name, 2, Offset::Offset(12));
    assert_eq!(position, consumer.position().unwrap());
}

#[test]
fn test_consumer_store_offset_commit() {
    let _r = env_logger::init();

    let topic_name = rand_test_topic();
    produce_messages(&topic_name, 10, &value_fn, &key_fn, Some(0), None);
    produce_messages(&topic_name, 11, &value_fn, &key_fn, Some(1), None);
    produce_messages(&topic_name, 12, &value_fn, &key_fn, Some(2), None);
    let mut config = HashMap::new();
    config.insert("enable.auto.offset.store", "false");
    let consumer = create_stream_consumer(&rand_test_group(), Some(config));
    consumer.subscribe(&[topic_name.as_str()]).unwrap();

    let _consumer_future = consumer.start()
        .take(33)
        .for_each(|message| {
            match message {
                Ok(m) => {
                    if m.partition() == 1 {
                        consumer.store_offset(&m).unwrap();
                    }
                },
                Err(e) => panic!("Error receiving message: {:?}", e)
            };
            Ok(())
        })
        .wait();

    // Commit the whole current state
    consumer.commit_consumer_state(CommitMode::Sync).unwrap();

    assert_eq!(consumer.fetch_watermarks(&topic_name, 0, 5000).unwrap(), (0, 10));
    assert_eq!(consumer.fetch_watermarks(&topic_name, 1, 5000).unwrap(), (0, 11));
    assert_eq!(consumer.fetch_watermarks(&topic_name, 2, 5000).unwrap(), (0, 12));

    let mut assignment = TopicPartitionList::new();
    assignment.add_partition_offset(&topic_name, 0, Offset::Invalid);
    assignment.add_partition_offset(&topic_name, 1, Offset::Invalid);
    assignment.add_partition_offset(&topic_name, 2, Offset::Invalid);
    assert_eq!(assignment, consumer.assignment().unwrap());

    let mut committed = TopicPartitionList::new();
    committed.add_partition_offset(&topic_name, 0, Offset::Invalid);
    committed.add_partition_offset(&topic_name, 1, Offset::Offset(11));
    committed.add_partition_offset(&topic_name, 2, Offset::Invalid);
    assert_eq!(committed, consumer.committed(5000).unwrap());

    let mut position = TopicPartitionList::new();
    position.add_partition_offset(&topic_name, 0, Offset::Offset(10));
    position.add_partition_offset(&topic_name, 1, Offset::Offset(11));
    position.add_partition_offset(&topic_name, 2, Offset::Offset(12));
    assert_eq!(position, consumer.position().unwrap());
}
<commit_msg>Add test for consumer error<commit_after>//! Test data production and consumption using high
//! level clients (StreamConsumer and FutureProducer).
extern crate env_logger;
extern crate futures;
extern crate rand;
extern crate rdkafka;

use futures::*;

use rdkafka::{Message, Timestamp};
use rdkafka::consumer::{Consumer, CommitMode};
use rdkafka::topic_partition_list::{Offset, TopicPartitionList};
use rdkafka::error::KafkaError;

mod utils;
use utils::*;

use std::time::{Duration, Instant};
use std::collections::HashMap;

// All messages should go to the same partition.
#[test]
fn test_produce_partition() {
    let _r = env_logger::init();

    let topic_name = rand_test_topic();
    let message_map = produce_messages(&topic_name, 100, &value_fn, &key_fn, Some(0), None);

    let res = message_map.iter()
        .filter(|&(&(partition, _), _)| partition == 0)
        .count();

    assert_eq!(res, 100);
}

// All produced messages should be consumed.
#[test]
fn test_produce_consume_base() {
    let _r = env_logger::init();

    let topic_name = rand_test_topic();
    let message_map = produce_messages(&topic_name, 100, &value_fn, &key_fn, None, None);
    let consumer = create_stream_consumer(&rand_test_group(), None);
    consumer.subscribe(&[topic_name.as_str()]).unwrap();

    let _consumer_future = consumer.start()
        .take(100)
        .for_each(|message| {
            match message {
                Ok(m) => {
                    let id = message_map[&(m.partition(), m.offset())];
                    match m.timestamp() {
                        Timestamp::CreateTime(timestamp) => assert!(timestamp > 1489495183000),
                        _ => panic!("Expected createtime for message timestamp")
                    };
                    assert_eq!(m.payload_view::<str>().unwrap().unwrap(), value_fn(id));
                    assert_eq!(m.key_view::<str>().unwrap().unwrap(), key_fn(id));
                    assert_eq!(m.topic(), topic_name.as_str());
                },
                Err(e) => panic!("Error receiving message: {:?}", e)
            };
            Ok(())
        })
        .wait();
}

// All produced messages should be consumed.
#[test]
fn test_produce_consume_base_assign() {
    let _r = env_logger::init();

    let topic_name = rand_test_topic();
    produce_messages(&topic_name, 10, &value_fn, &key_fn, Some(0), None);
    produce_messages(&topic_name, 10, &value_fn, &key_fn, Some(1), None);
    produce_messages(&topic_name, 10, &value_fn, &key_fn, Some(2), None);
    let consumer = create_stream_consumer(&rand_test_group(), None);
    let mut tpl = TopicPartitionList::new();
    tpl.add_partition_offset(&topic_name, 0, Offset::Beginning);
    tpl.add_partition_offset(&topic_name, 1, Offset::Offset(2));
    tpl.add_partition_offset(&topic_name, 2, Offset::Offset(9));
    consumer.assign(&tpl).unwrap();

    let mut partition_count = vec![0, 0, 0];

    let _consumer_future = consumer.start()
        .take(19)
        .for_each(|message| {
            match message {
                Ok(m) => partition_count[m.partition() as usize] += 1,
                Err(e) => panic!("Error receiving message: {:?}", e)
            };
            Ok(())
        })
        .wait();

    assert_eq!(partition_count, vec![10, 8, 1]);
}

// All produced messages should be consumed.
#[test]
fn test_produce_consume_with_timestamp() {
    let _r = env_logger::init();

    let topic_name = rand_test_topic();
    let message_map = produce_messages(&topic_name, 100, &value_fn, &key_fn, Some(0), Some(1111));
    let consumer = create_stream_consumer(&rand_test_group(), None);
    consumer.subscribe(&[topic_name.as_str()]).unwrap();

    let _consumer_future = consumer.start()
        .take(100)
        .for_each(|message| {
            match message {
                Ok(m) => {
                    let id = message_map[&(m.partition(), m.offset())];
                    assert_eq!(m.timestamp(), Timestamp::CreateTime(1111));
                    assert_eq!(m.payload_view::<str>().unwrap().unwrap(), value_fn(id));
                    assert_eq!(m.key_view::<str>().unwrap().unwrap(), key_fn(id));
                },
                Err(e) => panic!("Error receiving message: {:?}", e)
            };
            Ok(())
        })
        .wait();

    produce_messages(&topic_name, 10, &value_fn, &key_fn, Some(0), Some(999999));

    // Lookup the offsets
    let tpl = consumer.offsets_for_timestamp(999999, 10000).unwrap();
    let tp = tpl.find_partition(&topic_name, 0).unwrap();
    assert_eq!(tp.topic(), topic_name);
    assert_eq!(tp.offset(), Offset::Offset(100));
    assert_eq!(tp.partition(), 0);
    assert_eq!(tp.error(), Ok(()));
}

#[test]
fn test_consume_with_no_message_error() {
    let _r = env_logger::init();

    let consumer = create_stream_consumer(&rand_test_group(), None);

    let message_stream = consumer.start_with(Duration::from_millis(200), true);

    let mut first_poll_time = None;
    let mut timeouts_count = 0;
    for message in message_stream.wait() {
        match message {
            Ok(Err(KafkaError::NoMessageReceived)) => {
                // TODO: use entry interface for Options once available
                if first_poll_time.is_none() {
                    first_poll_time = Some(Instant::now());
                }
                timeouts_count += 1;
                if timeouts_count == 26 {
                    break;
                }
            }
            Ok(m) => panic!("A message was actually received: {:?}", m),
            Err(e) => panic!("Unexpected error while receiving message: {:?}", e)
        };
    }

    assert_eq!(timeouts_count, 26);
    // It should take 5000ms
    println!("Duration: {:?}", first_poll_time.unwrap().elapsed());
    assert!(first_poll_time.unwrap().elapsed() < Duration::from_millis(7000));
    assert!(first_poll_time.unwrap().elapsed() > Duration::from_millis(4500));
}

// TODO: add check that commit cb gets called correctly
#[test]
fn test_consumer_commit_message() {
    let _r = env_logger::init();

    let topic_name = rand_test_topic();
    produce_messages(&topic_name, 10, &value_fn, &key_fn, Some(0), None);
    produce_messages(&topic_name, 11, &value_fn, &key_fn, Some(1), None);
    produce_messages(&topic_name, 12, &value_fn, &key_fn, Some(2), None);
    let consumer = create_stream_consumer(&rand_test_group(), None);
    consumer.subscribe(&[topic_name.as_str()]).unwrap();

    let _consumer_future = consumer.start()
        .take(33)
        .for_each(|message| {
            match message {
                Ok(m) => {
                    if m.partition() == 1 {
                        consumer.commit_message(&m, CommitMode::Async).unwrap();
                    }
                },
                Err(e) => panic!("error receiving message: {:?}", e)
            };
            Ok(())
        })
        .wait();

    assert_eq!(consumer.fetch_watermarks(&topic_name, 0, 5000).unwrap(), (0, 10));
    assert_eq!(consumer.fetch_watermarks(&topic_name, 1, 5000).unwrap(), (0, 11));
    assert_eq!(consumer.fetch_watermarks(&topic_name, 2, 5000).unwrap(), (0, 12));

    let mut assignment = TopicPartitionList::new();
    assignment.add_partition_offset(&topic_name, 0, Offset::Invalid);
    assignment.add_partition_offset(&topic_name, 1, Offset::Invalid);
    assignment.add_partition_offset(&topic_name, 2, Offset::Invalid);
    assert_eq!(assignment, consumer.assignment().unwrap());

    let mut committed = TopicPartitionList::new();
    committed.add_partition_offset(&topic_name, 0, Offset::Invalid);
    committed.add_partition_offset(&topic_name, 1, Offset::Offset(11));
    committed.add_partition_offset(&topic_name, 2, Offset::Invalid);
    assert_eq!(committed, consumer.committed(5000).unwrap());

    let mut position = TopicPartitionList::new();
    position.add_partition_offset(&topic_name, 0, Offset::Offset(10));
    position.add_partition_offset(&topic_name, 1, Offset::Offset(11));
    position.add_partition_offset(&topic_name, 2, Offset::Offset(12));
    assert_eq!(position, consumer.position().unwrap());
}

#[test]
fn test_consumer_store_offset_commit() {
    let _r = env_logger::init();

    let topic_name = rand_test_topic();
    produce_messages(&topic_name, 10, &value_fn, &key_fn, Some(0), None);
    produce_messages(&topic_name, 11, &value_fn, &key_fn, Some(1), None);
    produce_messages(&topic_name, 12, &value_fn, &key_fn, Some(2), None);
    let mut config = HashMap::new();
    config.insert("enable.auto.offset.store", "false");
    config.insert("enable.partition.eof", "true");
    let consumer = create_stream_consumer(&rand_test_group(), Some(config));
    consumer.subscribe(&[topic_name.as_str()]).unwrap();

    let _consumer_future = consumer.start()
        .take(36)
        .for_each(|message| {
            match message {
                Ok(m) => {
                    if m.partition() == 1 {
                        consumer.store_offset(&m).unwrap();
                    }
                },
                Err(KafkaError::PartitionEOF(_)) => {},
                Err(e) => panic!("Error receiving message: {:?}", e)
            };
            Ok(())
        })
        .wait();

    // Commit the whole current state
    consumer.commit_consumer_state(CommitMode::Sync).unwrap();

    assert_eq!(consumer.fetch_watermarks(&topic_name, 0, 5000).unwrap(), (0, 10));
    assert_eq!(consumer.fetch_watermarks(&topic_name, 1, 5000).unwrap(), (0, 11));
    assert_eq!(consumer.fetch_watermarks(&topic_name, 2, 5000).unwrap(), (0, 12));

    let mut assignment = TopicPartitionList::new();
    assignment.add_partition_offset(&topic_name, 0, Offset::Invalid);
    assignment.add_partition_offset(&topic_name, 1, Offset::Invalid);
    assignment.add_partition_offset(&topic_name, 2, Offset::Invalid);
    assert_eq!(assignment, consumer.assignment().unwrap());

    let mut committed = TopicPartitionList::new();
    committed.add_partition_offset(&topic_name, 0, Offset::Invalid);
    committed.add_partition_offset(&topic_name, 1, Offset::Offset(11));
    committed.add_partition_offset(&topic_name, 2, Offset::Invalid);
    assert_eq!(committed, consumer.committed(5000).unwrap());

    let mut position = TopicPartitionList::new();
    position.add_partition_offset(&topic_name, 0, Offset::Offset(10));
    position.add_partition_offset(&topic_name, 1, Offset::Offset(11));
    position.add_partition_offset(&topic_name, 2, Offset::Offset(12));
    assert_eq!(position, consumer.position().unwrap());
}
<|endoftext|>


<commit_before>use std::sync::Arc;

use wayland_backend::server::{
    ClientData, ClientId, GlobalHandler, GlobalId, Handle, ObjectData, ObjectId,
};

use crate::{dispatch::ResourceData, Client, Dispatch, DisplayHandle, Resource};

pub(crate) struct GlobalData<I: Resource, D: GlobalDispatch<I>> {
    pub(crate) data: <D as GlobalDispatch<I>>::GlobalData,
}

impl<I: Resource + 'static, D: GlobalDispatch<I> + 'static> GlobalHandler<D> for GlobalData<I, D> {
    fn can_view(&self, id: ClientId, data: &Arc<dyn ClientData<D>>, _: GlobalId) -> bool {
        let client = Client { id, data: data.clone().into_any_arc() };
        <D as GlobalDispatch<I>>::can_view(client, &self.data)
    }

    fn bind(
        self: Arc<Self>,
        handle: &mut Handle<D>,
        data: &mut D,
        client_id: ClientId,
        _: GlobalId,
        object_id: ObjectId,
    ) -> Arc<dyn ObjectData<D>> {
        let mut handle = DisplayHandle::from_handle(handle);
        let client = Client::from_id(&mut handle, client_id).expect("Dead client in bind ?!");
        let resource = <I as Resource>::from_id(&mut handle, object_id)
            .expect("Wrong object_id in GlobalHandler ?!");

        let udata = data.bind(&mut handle, &client, &resource, &self.data);

        Arc::new(ResourceData::<I, _>::new(udata))
    }
}

pub trait GlobalDispatch<I: Resource>: Dispatch<I> {
    type GlobalData: Send + Sync + 'static;

    fn bind(
        &mut self,
        handle: &mut DisplayHandle<'_, Self>,
        client: &Client,
        resource: &I,
        global_data: &Self::GlobalData,
    ) -> <Self as Dispatch<I>>::UserData;

    fn can_view(_client: Client, _global_data: &Self::GlobalData) -> bool {
        true
    }
}
<commit_msg>server: global dispatch docs + delegate traits<commit_after>use std::sync::Arc;

use wayland_backend::server::{
    ClientData, ClientId, GlobalHandler, GlobalId, Handle, ObjectData, ObjectId,
};

use crate::{
    dispatch::{DelegateDispatch, DelegateDispatchBase, ResourceData},
    Client, Dispatch, DisplayHandle, Resource,
};

pub(crate) struct GlobalData<I: Resource, D: GlobalDispatch<I>> {
    pub(crate) data: <D as GlobalDispatch<I>>::GlobalData,
}

impl<I: Resource + 'static, D: GlobalDispatch<I> + 'static> GlobalHandler<D> for GlobalData<I, D> {
    fn can_view(&self, id: ClientId, data: &Arc<dyn ClientData<D>>, _: GlobalId) -> bool {
        let client = Client { id, data: data.clone().into_any_arc() };
        <D as GlobalDispatch<I>>::can_view(client, &self.data)
    }

    fn bind(
        self: Arc<Self>,
        handle: &mut Handle<D>,
        data: &mut D,
        client_id: ClientId,
        _: GlobalId,
        object_id: ObjectId,
    ) -> Arc<dyn ObjectData<D>> {
        let mut handle = DisplayHandle::from_handle(handle);
        let client = Client::from_id(&mut handle, client_id).expect("Dead client in bind ?!");
        let resource = <I as Resource>::from_id(&mut handle, object_id)
            .expect("Wrong object_id in GlobalHandler ?!");

        let udata = data.bind(&mut handle, &client, &resource, &self.data);

        Arc::new(ResourceData::<I, _>::new(udata))
    }
}

/// A trait which provides an implementation for handling advertisement of a global to clients with some type
/// of associated user data.
pub trait GlobalDispatch<I: Resource>: Dispatch<I> {
    /// Data associated with the global.
    type GlobalData: Send + Sync + 'static;

    /// Called when a client has bound this global.
    ///
    /// The return value of this function should contain user data to associate the object created by the
    /// client.
    fn bind(
        &mut self,
        handle: &mut DisplayHandle<'_, Self>,
        client: &Client,
        resource: &I,
        global_data: &Self::GlobalData,
    ) -> <Self as Dispatch<I>>::UserData;

    /// Checks if the global should be advertised to some client.
    ///
    /// The implementation of this function determines whether a client may see and bind some global. If this
    /// function returns false, the client will not be told the global exists and attempts to bind the global
    /// will raise a protocol error.
    ///
    /// One use of this function is implementing privileged protocols such as XWayland keyboard grabbing
    /// which must only be used by XWayland.
    ///
    /// The default implementation allows all clients to see the global.
    fn can_view(_client: Client, _global_data: &Self::GlobalData) -> bool {
        true
    }
}

/*
 * Dispatch delegation helpers
 */

/// The base trait used to define a delegate type to handle some global.
pub trait DelegateGlobalDispatchBase<I: Resource>: DelegateDispatchBase<I> {
    /// Data associated with the global.
    type GlobalData: Send + Sync + 'static;
}

/// A trait which defines a delegate type to handle some type of global.
///
/// This trait is useful for building modular handlers for globals.
pub trait DelegateGlobalDispatch<
    I: Resource,
    D: GlobalDispatch<I, GlobalData = <Self as DelegateGlobalDispatchBase<I>>::GlobalData>
        + Dispatch<I, UserData = Self::UserData>,
>: Sized + DelegateGlobalDispatchBase<I> + DelegateDispatch<I, D>
{
    /// Called when a client has bound this global.
    ///
    /// The return value of this function should contain user data to associate the object created by the
    /// client.
    fn bind(
        &mut self,
        handle: &mut DisplayHandle<'_, D>,
        client: &Client,
        resource: &I,
        global_data: &Self::GlobalData,
    ) -> D::UserData;

    /// Checks if the global should be advertised to some client.
    ///
    /// The implementation of this function determines whether a client may see and bind some global. If this
    /// function returns false, the client will not be told the global exists and attempts to bind the global
    /// will raise a protocol error.
    ///
    /// One use of this function is implementing privileged protocols such as XWayland keyboard grabbing
    /// which must only be used by XWayland.
    ///
    /// The default implementation allows all clients to see the global.
    fn can_view(_client: Client, _global_data: &Self::GlobalData) -> bool {
        true
    }
}
<|endoftext|>


<commit_before>use std::ops::Mul;
use std::marker::PhantomData;
use na::{Translate, Cross, Translation, Rotation};
use math::{Point, Vector, Isometry};
use geometry::shape::{Shape, Ball, Plane};
use geometry::query::algorithms::johnson_simplex::JohnsonSimplex;
use narrow_phase::{
    ContactDispatcher,
    ContactAlgorithm,
    BallBallContactGenerator,
    PlaneSupportMapContactGenerator,
    SupportMapPlaneContactGenerator,
    SupportMapSupportMapContactGenerator,
    CompositeShapeShapeContactGenerator,
    ShapeCompositeShapeContactGenerator,
    OneShotContactManifoldGenerator
};

/// Collision dispatcher for shapes defined by `ncollide_entities`.
pub struct DefaultContactDispatcher<P: Point, M> {
    _point_type:  PhantomData<P>,
    _matrix_type: PhantomData<M>,
}

impl<P: Point, M> DefaultContactDispatcher<P, M> {
    /// Creates a new basic collision dispatcher.
    pub fn new() -> DefaultContactDispatcher<P, M> {
        DefaultContactDispatcher {
            _point_type:  PhantomData,
            _matrix_type: PhantomData,
        }
    }
}

impl<P, M> ContactDispatcher<P, M> for DefaultContactDispatcher<P, M>
    where P: Point,
          P::Vect: Translate<P> + Cross,
          <P::Vect as Cross>::CrossProductType: Vector<Scalar = <P::Vect as Vector>::Scalar> +
                                                Mul<<P::Vect as Vector>::Scalar, Output = <P::Vect as Cross>::CrossProductType>, // FIXME: why do we need this?
          M: Isometry<P> + Translation<P::Vect> + Rotation<<P::Vect as Cross>::CrossProductType> {
    fn get_contact_algorithm(&self, a: &Shape<P, M>, b: &Shape<P, M>) -> Option<ContactAlgorithm<P, M>> {
        let a_is_ball = a.is_shape::<Ball<<P::Vect as Vector>::Scalar>>();
        let b_is_ball = b.is_shape::<Ball<<P::Vect as Vector>::Scalar>>();

        if a_is_ball && b_is_ball {
            Some(Box::new(BallBallContactGenerator::<P, M>::new()))
        }
        else if a.is_shape::<Plane<P::Vect>>() && b.is_support_map() {
            let wo_manifold = PlaneSupportMapContactGenerator::<P, M>::new();

            if !b_is_ball {
                let manifold = OneShotContactManifoldGenerator::new(wo_manifold);
                Some(Box::new(manifold))
            }
            else {
                Some(Box::new(wo_manifold))
            }
        }
        else if b.is_shape::<Plane<P::Vect>>() && a.is_support_map() {
            let wo_manifold = SupportMapPlaneContactGenerator::<P, M>::new();

            if !b_is_ball {
                let manifold = OneShotContactManifoldGenerator::new(wo_manifold);
                Some(Box::new(manifold))
            }
            else {
                Some(Box::new(wo_manifold))
            }
        }
        else if a.is_support_map() && b.is_support_map() {
            let simplex = JohnsonSimplex::new_w_tls();
            let wo_manifold = SupportMapSupportMapContactGenerator::new(simplex);

            if !b_is_ball {
                let manifold = OneShotContactManifoldGenerator::new(wo_manifold);
                Some(Box::new(manifold))
            }
            else {
                Some(Box::new(wo_manifold))
            }
        }
        else if a.is_composite_shape() {
            Some(Box::new(CompositeShapeShapeContactGenerator::<P, M>::new()))
        }
        else if b.is_composite_shape() {
            Some(Box::new(ShapeCompositeShapeContactGenerator::<P, M>::new()))
        }
        else {
            None
        }
    }
}
<commit_msg>Do not use contact manifold generation for balls.<commit_after>use std::ops::Mul;
use std::marker::PhantomData;
use na::{Translate, Cross, Translation, Rotation};
use math::{Point, Vector, Isometry};
use geometry::shape::{Shape, Ball, Plane};
use geometry::query::algorithms::johnson_simplex::JohnsonSimplex;
use narrow_phase::{
    ContactDispatcher,
    ContactAlgorithm,
    BallBallContactGenerator,
    PlaneSupportMapContactGenerator,
    SupportMapPlaneContactGenerator,
    SupportMapSupportMapContactGenerator,
    CompositeShapeShapeContactGenerator,
    ShapeCompositeShapeContactGenerator,
    OneShotContactManifoldGenerator
};

/// Collision dispatcher for shapes defined by `ncollide_entities`.
pub struct DefaultContactDispatcher<P: Point, M> {
    _point_type:  PhantomData<P>,
    _matrix_type: PhantomData<M>,
}

impl<P: Point, M> DefaultContactDispatcher<P, M> {
    /// Creates a new basic collision dispatcher.
    pub fn new() -> DefaultContactDispatcher<P, M> {
        DefaultContactDispatcher {
            _point_type:  PhantomData,
            _matrix_type: PhantomData,
        }
    }
}

impl<P, M> ContactDispatcher<P, M> for DefaultContactDispatcher<P, M>
    where P: Point,
          P::Vect: Translate<P> + Cross,
          <P::Vect as Cross>::CrossProductType: Vector<Scalar = <P::Vect as Vector>::Scalar> +
                                                Mul<<P::Vect as Vector>::Scalar, Output = <P::Vect as Cross>::CrossProductType>, // FIXME: why do we need this?
          M: Isometry<P> + Translation<P::Vect> + Rotation<<P::Vect as Cross>::CrossProductType> {
    fn get_contact_algorithm(&self, a: &Shape<P, M>, b: &Shape<P, M>) -> Option<ContactAlgorithm<P, M>> {
        let a_is_ball = a.is_shape::<Ball<<P::Vect as Vector>::Scalar>>();
        let b_is_ball = b.is_shape::<Ball<<P::Vect as Vector>::Scalar>>();

        if a_is_ball && b_is_ball {
            Some(Box::new(BallBallContactGenerator::<P, M>::new()))
        }
        else if a.is_shape::<Plane<P::Vect>>() && b.is_support_map() {
            let wo_manifold = PlaneSupportMapContactGenerator::<P, M>::new();

            if !b_is_ball {
                let manifold = OneShotContactManifoldGenerator::new(wo_manifold);
                Some(Box::new(manifold))
            }
            else {
                Some(Box::new(wo_manifold))
            }
        }
        else if b.is_shape::<Plane<P::Vect>>() && a.is_support_map() {
            let wo_manifold = SupportMapPlaneContactGenerator::<P, M>::new();

            if !a_is_ball {
                let manifold = OneShotContactManifoldGenerator::new(wo_manifold);
                Some(Box::new(manifold))
            }
            else {
                Some(Box::new(wo_manifold))
            }
        }
        else if a.is_support_map() && b.is_support_map() {
            let simplex     = JohnsonSimplex::new_w_tls();
            let wo_manifold = SupportMapSupportMapContactGenerator::new(simplex);

            if !a_is_ball && !b_is_ball {
                let manifold = OneShotContactManifoldGenerator::new(wo_manifold);
                Some(Box::new(manifold))
            }
            else {
                Some(Box::new(wo_manifold))
            }
        }
        else if a.is_composite_shape() {
            Some(Box::new(CompositeShapeShapeContactGenerator::<P, M>::new()))
        }
        else if b.is_composite_shape() {
            Some(Box::new(ShapeCompositeShapeContactGenerator::<P, M>::new()))
        }
        else {
            None
        }
    }
}
<|endoftext|>


<commit_before>use image::ImageResult;
use mli::{Backward, Forward, Graph, Train};
use mli_conv::Conv2;
use mli_ndarray::Map2One;
use mli_relu::Blu;
use ndarray::{array, s, Array, Array2, OwnedRepr};
use ndarray_image::{open_gray_image, save_gray_image};
use rand_core::SeedableRng;
use rand_distr::{Distribution, Normal};
use rand_pcg::Pcg64;
use std::path::{Path, PathBuf};
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
#[structopt(name = "image", about = "Loads and convolves an image as a test")]
struct Opt {
    /// Number of epochs
    #[structopt(short = "t", default_value = "2000")]
    epochs: usize,
    /// Number of epochs per output image
    #[structopt(short = "s", default_value = "1")]
    show_every: usize,
    /// Initial learning rate
    #[structopt(short = "i", default_value = "0.000000001")]
    initial_learning_rate: f32,
    /// Learning rate multiplier per epoch
    #[structopt(short = "m", default_value = "1.003")]
    learning_rate_multiplier: f32,
    /// Seed
    #[structopt(short = "z", default_value = "42")]
    seed: u32,
    /// Beta value for NAG
    #[structopt(short = "b", default_value = "0.9")]
    beta: f32,
    /// File to load
    #[structopt(parse(from_os_str))]
    file: PathBuf,
    /// Output directory
    #[structopt(parse(from_os_str))]
    output_dir: PathBuf,
}

fn sobel(image: &Array2<f32>) -> Array2<f32> {
    let down_filter = array![[-1.0, -2.0, -1.0], [0.0, 0.0, 0.0], [1.0, 2.0, 1.0],];
    let right_filter = down_filter.t().to_owned();
    let down = Conv2::new(down_filter).forward(&image).1;
    let right = Conv2::new(right_filter).forward(&image).1;
    (down.map(|f| f.powi(2)) + right.map(|f| f.powi(2))).map(|f| f.sqrt())
}

fn open_image(path: impl AsRef<Path>) -> ImageResult<Array2<f32>> {
    let image = open_gray_image(path)?;
    Ok(image.map(|&n| f32::from(n) / 255.0))
}

fn save_image(path: impl AsRef<Path>, image: &Array2<f32>) -> ImageResult<()> {
    let image = image.map(|&n| num::clamp(n * 255.0, 0.0, 255.0) as u8);
    save_gray_image(path, image.view())
}

fn main() -> ImageResult<()> {
    let opt = Opt::from_args();
    let image = open_image(opt.file)?;
    let sobel_image = sobel(&image);
    let conv_layers = 5;
    let filter_radius = 1usize;
    let filter_len = (filter_radius * 2 + 1).pow(2);
    let padding = (conv_layers * filter_radius - 1) as i32;
    #[allow(clippy::deref_addrof)]
    let expected = sobel_image.slice(s![padding..-padding, padding..-padding]);
    save_image(opt.output_dir.join("actual.png"), &expected.to_owned())?;
    let mut prng = Pcg64::from_seed([
        (opt.seed & 0xFF) as u8,
        ((opt.seed >> 8) & 0xFF) as u8,
        ((opt.seed >> 16) & 0xFF) as u8,
        ((opt.seed >> 24) & 0xFF) as u8,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
    ]);
    let mut random_filter = |mean: f32, variance: f32| -> Conv2<OwnedRepr<f32>> {
        // Xavier initialize by changing the variance to be 1/N where N is the number of neurons.
        Conv2::new(
            Array::from_iter(
                Normal::new(mean / filter_len as f32, variance / filter_len as f32)
                    .unwrap()
                    .sample_iter(&mut prng)
                    .take(filter_len),
            )
            .into_shape((filter_radius * 2 + 1, filter_radius * 2 + 1))
            .unwrap(),
        )
    };
    let mut prng = Pcg64::from_seed([
        (opt.seed & 0xFF) as u8,
        ((opt.seed >> 8) & 0xFF) as u8,
        ((opt.seed >> 16) & 0xFF) as u8,
        ((opt.seed >> 24) & 0xFF) as u8,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
    ]);
    let mut random_blu = |mean: f32, variance: f32| -> Blu {
        // Xavier initialize by changing the variance to be 1/N where N is the number of neurons.
        let distr = Normal::new(mean, variance).unwrap();
        Blu::new(distr.sample(&mut prng), distr.sample(&mut prng))
    };
    let mut generate_filter = || {
        random_filter(0.0, 10.0f32.powi(0))
            .chain(Map2One(random_blu(0.0, 0.05)))
            .chain(random_filter(0.0, 10.0f32.powi(0)))
            .chain(Map2One(random_blu(0.0, 0.05)))
            .chain(random_filter(0.0, 10.0f32.powi(0)))
            .chain(Map2One(random_blu(0.0, 0.05)))
            .chain(random_filter(0.0, 10.0f32.powi(0)))
            .chain(Map2One(random_blu(0.0, 0.05)))
            .chain(random_filter(64.0, 10.0f32.powi(3)))
    };

    loop {
        let mut train_filter = generate_filter();
        let mut learn_rate = opt.initial_learning_rate;
        // Weird hack to initialize the momentum to zero without knowing the shape of the ndarray in advance.
        let mut momentum = {
            let (internal, output) = train_filter.forward(&image);
            train_filter.backward_train(&image, &internal, &output)
        };
        momentum *= 0.0f32;
        for i in 0..opt.epochs {
            // Compute beta * momentum.
            momentum *= opt.beta;
            train_filter.train(&momentum);
            let (internal, output) = train_filter.forward(&image);
            // Show the image if the frame is divisible by show_every.
            if i % opt.show_every == 0 {
                save_image(opt.output_dir.join(format!("epoch{:04}.png", i)), &output)?;
            }
            // Compute the loss for display only (we don't actually need the loss itself for backprop, just its derivative).
            let loss = (output.clone() - expected.view()).map(|n| n.powi(2)).sum();
            eprintln!("epoch {:04} loss: {}, learn_rate: {}", i, loss, learn_rate);
            if i > 10 && loss > 80000.0 {
                eprintln!("loss > 80000.0 at epoch {}; starting over", i);
                break;
            }
            if !loss.is_normal() {
                eprintln!("abnormal loss at epoch {}; starting over", i);
                break;
            }
            // The loss function is (n - t)^2, so 2*(n - t) is dE/df where
            // `E` is loss and `f` is output.
            let delta_loss = (output - expected).map(|n| 2.0 * n);
            // Compute the output delta.
            let output_delta = -learn_rate * delta_loss;
            // Compute the trainable variable delta.
            let mut train_delta = train_filter.backward_train(&image, &internal, &output_delta);
            // Make the train delta a small component.
            train_delta *= 1.0 - opt.beta;
            train_filter.train(&train_delta);
            // Add the small component to the momentum.
            momentum += train_delta;
            learn_rate *= opt.learning_rate_multiplier;
        }
    }
}
<commit_msg>add more loss checkpoints<commit_after>use image::ImageResult;
use mli::{Backward, Forward, Graph, Train};
use mli_conv::Conv2;
use mli_ndarray::Map2One;
use mli_relu::Blu;
use ndarray::{array, s, Array, Array2, OwnedRepr};
use ndarray_image::{open_gray_image, save_gray_image};
use rand_core::SeedableRng;
use rand_distr::{Distribution, Normal};
use rand_pcg::Pcg64;
use std::path::{Path, PathBuf};
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
#[structopt(name = "image", about = "Loads and convolves an image as a test")]
struct Opt {
    /// Number of epochs
    #[structopt(short = "t", default_value = "2000")]
    epochs: usize,
    /// Number of epochs per output image
    #[structopt(short = "s", default_value = "1")]
    show_every: usize,
    /// Initial learning rate
    #[structopt(short = "i", default_value = "0.000000001")]
    initial_learning_rate: f32,
    /// Learning rate multiplier per epoch
    #[structopt(short = "m", default_value = "1.003")]
    learning_rate_multiplier: f32,
    /// Seed
    #[structopt(short = "z", default_value = "42")]
    seed: u32,
    /// Beta value for NAG
    #[structopt(short = "b", default_value = "0.9")]
    beta: f32,
    /// File to load
    #[structopt(parse(from_os_str))]
    file: PathBuf,
    /// Output directory
    #[structopt(parse(from_os_str))]
    output_dir: PathBuf,
}

fn sobel(image: &Array2<f32>) -> Array2<f32> {
    let down_filter = array![[-1.0, -2.0, -1.0], [0.0, 0.0, 0.0], [1.0, 2.0, 1.0],];
    let right_filter = down_filter.t().to_owned();
    let down = Conv2::new(down_filter).forward(&image).1;
    let right = Conv2::new(right_filter).forward(&image).1;
    (down.map(|f| f.powi(2)) + right.map(|f| f.powi(2))).map(|f| f.sqrt())
}

fn open_image(path: impl AsRef<Path>) -> ImageResult<Array2<f32>> {
    let image = open_gray_image(path)?;
    Ok(image.map(|&n| f32::from(n) / 255.0))
}

fn save_image(path: impl AsRef<Path>, image: &Array2<f32>) -> ImageResult<()> {
    let image = image.map(|&n| num::clamp(n * 255.0, 0.0, 255.0) as u8);
    save_gray_image(path, image.view())
}

fn main() -> ImageResult<()> {
    let opt = Opt::from_args();
    let image = open_image(opt.file)?;
    let sobel_image = sobel(&image);
    let conv_layers = 5;
    let filter_radius = 1usize;
    let filter_len = (filter_radius * 2 + 1).pow(2);
    let padding = (conv_layers * filter_radius - 1) as i32;
    #[allow(clippy::deref_addrof)]
    let expected = sobel_image.slice(s![padding..-padding, padding..-padding]);
    save_image(opt.output_dir.join("actual.png"), &expected.to_owned())?;
    let mut prng = Pcg64::from_seed([
        (opt.seed & 0xFF) as u8,
        ((opt.seed >> 8) & 0xFF) as u8,
        ((opt.seed >> 16) & 0xFF) as u8,
        ((opt.seed >> 24) & 0xFF) as u8,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
    ]);
    let mut random_filter = |mean: f32, variance: f32| -> Conv2<OwnedRepr<f32>> {
        // Xavier initialize by changing the variance to be 1/N where N is the number of neurons.
        Conv2::new(
            Array::from_iter(
                Normal::new(mean / filter_len as f32, variance / filter_len as f32)
                    .unwrap()
                    .sample_iter(&mut prng)
                    .take(filter_len),
            )
            .into_shape((filter_radius * 2 + 1, filter_radius * 2 + 1))
            .unwrap(),
        )
    };
    let mut prng = Pcg64::from_seed([
        (opt.seed & 0xFF) as u8,
        ((opt.seed >> 8) & 0xFF) as u8,
        ((opt.seed >> 16) & 0xFF) as u8,
        ((opt.seed >> 24) & 0xFF) as u8,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
    ]);
    let mut random_blu = |mean: f32, variance: f32| -> Blu {
        // Xavier initialize by changing the variance to be 1/N where N is the number of neurons.
        let distr = Normal::new(mean, variance).unwrap();
        Blu::new(distr.sample(&mut prng), distr.sample(&mut prng))
    };
    let mut generate_filter = || {
        random_filter(0.0, 10.0f32.powi(0))
            .chain(Map2One(random_blu(0.0, 0.05)))
            .chain(random_filter(0.0, 10.0f32.powi(0)))
            .chain(Map2One(random_blu(0.0, 0.05)))
            .chain(random_filter(0.0, 10.0f32.powi(0)))
            .chain(Map2One(random_blu(0.0, 0.05)))
            .chain(random_filter(0.0, 10.0f32.powi(0)))
            .chain(Map2One(random_blu(0.0, 0.05)))
            .chain(random_filter(64.0, 10.0f32.powi(3)))
    };

    loop {
        let mut train_filter = generate_filter();
        let mut learn_rate = opt.initial_learning_rate;
        // Weird hack to initialize the momentum to zero without knowing the shape of the ndarray in advance.
        let mut momentum = {
            let (internal, output) = train_filter.forward(&image);
            train_filter.backward_train(&image, &internal, &output)
        };
        momentum *= 0.0f32;
        for i in 0..opt.epochs {
            // Compute beta * momentum.
            momentum *= opt.beta;
            train_filter.train(&momentum);
            let (internal, output) = train_filter.forward(&image);
            // Show the image if the frame is divisible by show_every.
            if i % opt.show_every == 0 {
                save_image(opt.output_dir.join(format!("epoch{:04}.png", i)), &output)?;
            }
            // Compute the loss for display only (we don't actually need the loss itself for backprop, just its derivative).
            let loss = (output.clone() - expected.view()).map(|n| n.powi(2)).sum();
            eprintln!("epoch {:04} loss: {}, learn_rate: {}", i, loss, learn_rate);
            if i > 15 && loss > 80000.0 {
                eprintln!("loss > 80000.0 at epoch {}; starting over", i);
                break;
            }
            if i > 50 && loss > 70000.0 {
                eprintln!("loss > 70000.0 at epoch {}; starting over", i);
                break;
            }
            if i > 150 && loss > 60000.0 {
                eprintln!("loss > 60000.0 at epoch {}; starting over", i);
                break;
            }
            if i > 500 && loss > 30000.0 {
                eprintln!("loss > 60000.0 at epoch {}; starting over", i);
                break;
            }
            if !loss.is_normal() {
                eprintln!("abnormal loss at epoch {}; starting over", i);
                break;
            }
            // The loss function is (n - t)^2, so 2*(n - t) is dE/df where
            // `E` is loss and `f` is output.
            let delta_loss = (output - expected).map(|n| 2.0 * n);
            // Compute the output delta.
            let output_delta = -learn_rate * delta_loss;
            // Compute the trainable variable delta.
            let mut train_delta = train_filter.backward_train(&image, &internal, &output_delta);
            // Make the train delta a small component.
            train_delta *= 1.0 - opt.beta;
            train_filter.train(&train_delta);
            // Add the small component to the momentum.
            momentum += train_delta;
            learn_rate *= opt.learning_rate_multiplier;
        }
    }
}
<|endoftext|>


<commit_before>#![feature(plugin_registrar, rustc_private)]

extern crate itertools;
extern crate rustc;
extern crate rustc_plugin;
extern crate syntax;

use itertools::Itertools;
use rustc_plugin::Registry;
use syntax::ast::{self, Ident, TraitRef, Ty, TyKind};
use syntax::ast::LitKind::Str;
use syntax::ast::MetaItemKind::NameValue;
use syntax::codemap::Spanned;
use syntax::ext::base::{ExtCtxt, MacResult, DummyResult, MacEager};
use syntax::ext::quote::rt::Span;
use syntax::parse::{self, token, PResult};
use syntax::parse::parser::{Parser, PathStyle};
use syntax::parse::token::intern_and_get_ident;
use syntax::ptr::P;
use syntax::tokenstream::TokenTree;
use syntax::util::small_vector::SmallVector;

fn snake_to_camel(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree]) -> Box<MacResult + 'static> {
    let mut parser = parse::new_parser_from_tts(cx.parse_sess(), cx.cfg(), tts.into());
    // The `expand_expr` method is called so that any macro calls in the
    // parsed expression are expanded.

    let mut item = match parser.parse_trait_item() {
        Ok(s) => s,
        Err(mut diagnostic) => {
            diagnostic.emit();
            return DummyResult::any(sp);
        }
    };

    if let Err(mut diagnostic) = parser.expect(&token::Eof) {
        diagnostic.emit();
        return DummyResult::any(sp);
    }

    let old_ident = convert(&mut item.ident);

    // As far as I know, it's not possible in macro_rules! to reference an $ident in a doc string,
    // so this is the hacky workaround.
    //
    // This code looks intimidating, but it's just iterating through the trait item's attributes
    // (NameValues), filtering out non-doc attributes, and replacing any {} in the doc string with
    // the original, snake_case ident.
    for meta_item in item.attrs.iter_mut().map(|attr| &mut attr.node.value) {
        let updated = match meta_item.node {
            NameValue(ref name, _) if name == "doc" => {
                let mut updated = (**meta_item).clone();
                if let NameValue(_, Spanned { node: Str(ref mut doc, _), .. }) = updated.node {
                    let updated_doc = doc.replace("{}", &old_ident);
                    *doc = intern_and_get_ident(&updated_doc);
                } else {
                    unreachable!()
                };
                Some(P(updated))
            }
            _ => None,
        };
        if let Some(updated) = updated {
            *meta_item = updated;
        }
    }

    MacEager::trait_items(SmallVector::one(item))
}

fn impl_snake_to_camel(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree]) -> Box<MacResult + 'static> {
    let mut parser = parse::new_parser_from_tts(cx.parse_sess(), cx.cfg(), tts.into());
    // The `expand_expr` method is called so that any macro calls in the
    // parsed expression are expanded.

    let mut item = match parser.parse_impl_item() {
        Ok(s) => s,
        Err(mut diagnostic) => {
            diagnostic.emit();
            return DummyResult::any(sp);
        }
    };

    if let Err(mut diagnostic) = parser.expect(&token::Eof) {
        diagnostic.emit();
        return DummyResult::any(sp);
    }

    convert(&mut item.ident);
    MacEager::impl_items(SmallVector::one(item))
}

fn ty_snake_to_camel(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree]) -> Box<MacResult + 'static> {
    let mut parser = parse::new_parser_from_tts(cx.parse_sess(), cx.cfg(), tts.into());
    // The `expand_expr` method is called so that any macro calls in the
    // parsed expression are expanded.

    let mut ty = match parser.parse_ty_path() {
        Ok(s) => s,
        Err(mut diagnostic) => {
            diagnostic.emit();
            return DummyResult::any(sp);
        }
    };

    if let Err(mut diagnostic) = parser.expect(&token::Eof) {
        diagnostic.emit();
        return DummyResult::any(sp);
    }

    // Only capitalize the final segment
    if let TyKind::Path(_, ref mut path) = ty {
        convert(&mut path.segments.last_mut().unwrap().identifier);
    } else {
        unreachable!()
    }
    MacEager::ty(P(Ty {
        id: ast::DUMMY_NODE_ID,
        node: ty,
        span: sp,
    }))
}

/// Converts an ident in-place to CamelCase and returns the previous ident.
fn convert(ident: &mut Ident) -> String {
    let ident_str = ident.to_string();
    let mut camel_ty = String::new();

    {
        // Find the first non-underscore and add it capitalized.
        let mut chars = ident_str.chars();

        // Find the first non-underscore char, uppercase it, and append it.
        // Guaranteed to succeed because all idents must have at least one non-underscore char.
        camel_ty.extend(chars.find(|&c| c != '_').unwrap().to_uppercase());

        // When we find an underscore, we remove it and capitalize the next char. To do this,
        // we need to ensure the next char is not another underscore.
        let mut chars = chars.coalesce(|c1, c2| {
            if c1 == '_' && c2 == '_' {
                Ok(c1)
            } else {
                Err((c1, c2))
            }
        });

        while let Some(c) = chars.next() {
            if c != '_' {
                camel_ty.push(c);
            } else {
                if let Some(c) = chars.next() {
                    camel_ty.extend(c.to_uppercase());
                }
            }
        }
    }

    // The Fut suffix is hardcoded right now; this macro isn't really meant to be general-purpose.
    camel_ty.push_str("Fut");

    *ident = Ident::with_empty_ctxt(token::intern(&camel_ty));
    ident_str
}

trait ParseTraitRef {
    fn parse_trait_ref(&mut self) -> PResult<TraitRef>;
}

impl<'a> ParseTraitRef for Parser<'a> {
    /// Parse a::B<String,i32>
    fn parse_trait_ref(&mut self) -> PResult<TraitRef> {
        Ok(TraitRef {
            path: self.parse_path(PathStyle::Type)?,
            ref_id: ast::DUMMY_NODE_ID,
        })
    }
}

#[plugin_registrar]
#[doc(hidden)]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_macro("snake_to_camel", snake_to_camel);
    reg.register_macro("impl_snake_to_camel", impl_snake_to_camel);
    reg.register_macro("ty_snake_to_camel", ty_snake_to_camel);
}
<commit_msg>Track nightly compiler<commit_after>#![feature(plugin_registrar, rustc_private)]

extern crate itertools;
extern crate rustc;
extern crate rustc_plugin;
extern crate syntax;

use itertools::Itertools;
use rustc_plugin::Registry;
use syntax::ast::{self, Ident, TraitRef, Ty, TyKind};
use syntax::ast::LitKind::Str;
use syntax::ast::MetaItemKind::NameValue;
use syntax::codemap::Spanned;
use syntax::ext::base::{ExtCtxt, MacResult, DummyResult, MacEager};
use syntax::ext::quote::rt::Span;
use syntax::parse::{self, token, PResult};
use syntax::parse::parser::{Parser, PathStyle};
use syntax::parse::token::intern_and_get_ident;
use syntax::ptr::P;
use syntax::tokenstream::TokenTree;
use syntax::util::small_vector::SmallVector;

fn snake_to_camel(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree]) -> Box<MacResult + 'static> {
    let mut parser = parse::new_parser_from_tts(cx.parse_sess(), cx.cfg().clone(), tts.into());
    // The `expand_expr` method is called so that any macro calls in the
    // parsed expression are expanded.

    let mut item = match parser.parse_trait_item() {
        Ok(s) => s,
        Err(mut diagnostic) => {
            diagnostic.emit();
            return DummyResult::any(sp);
        }
    };

    if let Err(mut diagnostic) = parser.expect(&token::Eof) {
        diagnostic.emit();
        return DummyResult::any(sp);
    }

    let old_ident = convert(&mut item.ident);

    // As far as I know, it's not possible in macro_rules! to reference an $ident in a doc string,
    // so this is the hacky workaround.
    //
    // This code looks intimidating, but it's just iterating through the trait item's attributes
    // (NameValues), filtering out non-doc attributes, and replacing any {} in the doc string with
    // the original, snake_case ident.
    for meta_item in item.attrs.iter_mut().map(|attr| &mut attr.node.value) {
        let updated = match meta_item.node {
            NameValue(ref name, _) if name == "doc" => {
                let mut updated = (**meta_item).clone();
                if let NameValue(_, Spanned { node: Str(ref mut doc, _), .. }) = updated.node {
                    let updated_doc = doc.replace("{}", &old_ident);
                    *doc = intern_and_get_ident(&updated_doc);
                } else {
                    unreachable!()
                };
                Some(P(updated))
            }
            _ => None,
        };
        if let Some(updated) = updated {
            *meta_item = updated;
        }
    }

    MacEager::trait_items(SmallVector::one(item))
}

fn impl_snake_to_camel(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree]) -> Box<MacResult + 'static> {
    let mut parser = parse::new_parser_from_tts(cx.parse_sess(), cx.cfg().clone(), tts.into());
    // The `expand_expr` method is called so that any macro calls in the
    // parsed expression are expanded.

    let mut item = match parser.parse_impl_item() {
        Ok(s) => s,
        Err(mut diagnostic) => {
            diagnostic.emit();
            return DummyResult::any(sp);
        }
    };

    if let Err(mut diagnostic) = parser.expect(&token::Eof) {
        diagnostic.emit();
        return DummyResult::any(sp);
    }

    convert(&mut item.ident);
    MacEager::impl_items(SmallVector::one(item))
}

fn ty_snake_to_camel(cx: &mut ExtCtxt, sp: Span, tts: &[TokenTree]) -> Box<MacResult + 'static> {
    let mut parser = parse::new_parser_from_tts(cx.parse_sess(), cx.cfg().clone(), tts.into());
    // The `expand_expr` method is called so that any macro calls in the
    // parsed expression are expanded.

    let mut ty = match parser.parse_ty_path() {
        Ok(s) => s,
        Err(mut diagnostic) => {
            diagnostic.emit();
            return DummyResult::any(sp);
        }
    };

    if let Err(mut diagnostic) = parser.expect(&token::Eof) {
        diagnostic.emit();
        return DummyResult::any(sp);
    }

    // Only capitalize the final segment
    if let TyKind::Path(_, ref mut path) = ty {
        convert(&mut path.segments.last_mut().unwrap().identifier);
    } else {
        unreachable!()
    }
    MacEager::ty(P(Ty {
        id: ast::DUMMY_NODE_ID,
        node: ty,
        span: sp,
    }))
}

/// Converts an ident in-place to CamelCase and returns the previous ident.
fn convert(ident: &mut Ident) -> String {
    let ident_str = ident.to_string();
    let mut camel_ty = String::new();

    {
        // Find the first non-underscore and add it capitalized.
        let mut chars = ident_str.chars();

        // Find the first non-underscore char, uppercase it, and append it.
        // Guaranteed to succeed because all idents must have at least one non-underscore char.
        camel_ty.extend(chars.find(|&c| c != '_').unwrap().to_uppercase());

        // When we find an underscore, we remove it and capitalize the next char. To do this,
        // we need to ensure the next char is not another underscore.
        let mut chars = chars.coalesce(|c1, c2| {
            if c1 == '_' && c2 == '_' {
                Ok(c1)
            } else {
                Err((c1, c2))
            }
        });

        while let Some(c) = chars.next() {
            if c != '_' {
                camel_ty.push(c);
            } else {
                if let Some(c) = chars.next() {
                    camel_ty.extend(c.to_uppercase());
                }
            }
        }
    }

    // The Fut suffix is hardcoded right now; this macro isn't really meant to be general-purpose.
    camel_ty.push_str("Fut");

    *ident = Ident::with_empty_ctxt(token::intern(&camel_ty));
    ident_str
}

trait ParseTraitRef {
    fn parse_trait_ref(&mut self) -> PResult<TraitRef>;
}

impl<'a> ParseTraitRef for Parser<'a> {
    /// Parse a::B<String,i32>
    fn parse_trait_ref(&mut self) -> PResult<TraitRef> {
        Ok(TraitRef {
            path: self.parse_path(PathStyle::Type)?,
            ref_id: ast::DUMMY_NODE_ID,
        })
    }
}

#[plugin_registrar]
#[doc(hidden)]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_macro("snake_to_camel", snake_to_camel);
    reg.register_macro("impl_snake_to_camel", impl_snake_to_camel);
    reg.register_macro("ty_snake_to_camel", ty_snake_to_camel);
}
<|endoftext|>


<commit_before>// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use std::io;
use std::ops::Deref;
use std::os::unix::io::{AsRawFd, RawFd};

use libc::{
    epoll_create1, epoll_ctl, epoll_event, epoll_wait, EPOLLERR, EPOLLET, EPOLLEXCLUSIVE, EPOLLHUP,
    EPOLLIN, EPOLLONESHOT, EPOLLOUT, EPOLLPRI, EPOLLRDHUP, EPOLLWAKEUP, EPOLL_CLOEXEC,
    EPOLL_CTL_ADD, EPOLL_CTL_DEL, EPOLL_CTL_MOD,
};

use utils::syscall::SyscallReturnCode;

/// Wrapper over EPOLL_CTL_* operations that can be performed on a file descriptor.
#[repr(i32)]
pub enum ControlOperation {
    /// Add a file descriptor to the interest list.
    Add = EPOLL_CTL_ADD,
    /// Change the settings associated with a file descriptor that is
    /// already in the interest list.
    Modify = EPOLL_CTL_MOD,
    /// Remove a file descriptor from the interest list.
    Delete = EPOLL_CTL_DEL,
}

bitflags! {
    /// The type of events we can monitor a file descriptor for.
    pub struct EventSet: u32 {
        /// The associated file descriptor is available for read operations.
        const IN = EPOLLIN as u32;
        /// The associated file descriptor is available for write operations.
        const OUT = EPOLLOUT as u32;
        /// Error condition happened on the associated file descriptor.
        const ERROR = EPOLLERR as u32;
        /// This can be used to detect peer shutdown when using Edge Triggered monitoring.
        const READ_HANG_UP = EPOLLRDHUP as u32;
        /// Sets the Edge Triggered behavior for the associated file descriptor.
        /// The default behavior is Level Triggered.
        const EDGE_TRIGGERED = EPOLLET as u32;
        /// Hang up happened on the associated file descriptor. Note that `epoll_wait`
        /// will always wait for this event and it is not necessary to set it in events.
        const HANG_UP = EPOLLHUP as u32;
        /// There is an exceptional condition on that file descriptor. It is mostly used to
        /// set high priority for some data.
        const PRIORITY = EPOLLPRI as u32;
        /// The event is considered as being "processed" from the time when it is returned
        /// by a call to `epoll_wait` until the next call to `epoll_wait` on the same
        /// epoll file descriptor, the closure of that file descriptor, the removal of the
        /// event file descriptor via EPOLL_CTL_DEL, or the clearing of EPOLLWAKEUP
        /// for the event file descriptor via EPOLL_CTL_MOD.
        const WAKE_UP = EPOLLWAKEUP as u32;
        /// Sets the one-shot behavior for the associated file descriptor.
        const ONE_SHOT = EPOLLONESHOT as u32;
        /// Sets an exclusive wake up mode for the epoll file descriptor that is being
        /// attached to the associated file descriptor.
        /// When a wake up event occurs and multiple epoll file descriptors are attached to
        /// the same target file using this mode, one or more of the epoll file descriptors
        /// will receive an event with `epoll_wait`. The default here is for all those file
        /// descriptors to receive an event.
        const EXCLUSIVE = EPOLLEXCLUSIVE as u32;
    }
}

/// Wrapper over 'libc::epoll_event'.
///
// We are using `transparent` here to be super sure that this struct and its fields
// have the same alignment as those from the `epoll_event` struct from C.
#[repr(transparent)]
#[derive(Clone)]
pub struct EpollEvent(epoll_event);

impl Deref for EpollEvent {
    type Target = epoll_event;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl Default for EpollEvent {
    fn default() -> Self {
        EpollEvent(epoll_event {
            events: 0u32,
            u64: 0u64,
        })
    }
}

impl EpollEvent {
    /// Create a new epoll_event instance with the following fields: `events`, which contains
    /// an event mask and `data` which represents a user data variable. `data` field can be
    /// a fd on which we want to monitor the events specified by `events`.
    pub fn new(events: EventSet, data: u64) -> Self {
        EpollEvent(epoll_event {
            events: events.bits(),
            u64: data,
        })
    }

    /// Returns the `events` from `libc::epoll_event`.
    pub fn events(&self) -> u32 {
        self.events
    }

    /// Returns the `EventSet` corresponding to `epoll_event.events`.
    ///
    /// # Panics
    ///
    /// Panics if `libc::epoll_event` contains invalid events.
    pub fn event_set(&self) -> EventSet {
        // This unwrap is safe because `epoll_events` can only be user created or
        // initialized by the kernel.We trust the kernel to only send us valid
        // events. The user can only initialize `epoll_events` using valid events.
        EventSet::from_bits(self.events()).unwrap()
    }

    /// Returns the `data` from the `libc::epoll_event`.
    pub fn data(&self) -> u64 {
        self.u64
    }

    /// Converts the `libc::epoll_event` data to a RawFd.
    ///
    /// This conversion is lossy when the data does not correspond to a RawFd
    /// (data does not fit in a i32).
    pub fn fd(&self) -> RawFd {
        self.u64 as i32
    }
}

/// Wrapper over epoll functionality.
#[derive(Clone, Debug)]
pub struct Epoll {
    epoll_fd: RawFd,
}

impl Epoll {
    /// Create a new epoll file descriptor.
    pub fn new() -> io::Result<Self> {
        let epoll_fd = SyscallReturnCode(unsafe { epoll_create1(EPOLL_CLOEXEC) }).into_result()?;
        Ok(Epoll { epoll_fd })
    }

    /// Wrapper for `libc::epoll_ctl`.
    ///
    /// This can be used for adding, modifying or removing a file descriptor in the
    /// interest list of the epoll instance.
    ///
    /// # Arguments
    ///
    /// * `operation` refers to the action to be performed on the file descriptor.
    /// * `fd` is the file descriptor on which we want to perform `operation`.
    /// * `event` refers to the `epoll_event` instance that is linked to `fd`.
    pub fn ctl(
        &self,
        operation: ControlOperation,
        fd: RawFd,
        event: &EpollEvent,
    ) -> io::Result<()> {
        // Safe because we give a valid epoll file descriptor, a valid file descriptor to watch,
        // as well as a valid epoll_event structure. We also check the return value.
        SyscallReturnCode(unsafe {
            epoll_ctl(
                self.epoll_fd,
                operation as i32,
                fd,
                event as *const EpollEvent as *mut epoll_event,
            )
        })
        .into_empty_result()
    }

    /// Wrapper for `libc::epoll_wait`.
    /// Returns the number of file descriptors in the interest list that became ready
    /// for I/O or `errno` if an error occurred.
    ///
    /// # Arguments
    ///
    /// * `max_events` is the maximum number of events that we want to be returned in
    /// `events` buffer.
    /// * `timeout` specifies for how long the `epoll_wait` system call will block
    /// (measured in milliseconds).
    /// * `events` points to a memory area that will be used for storing the events
    /// returned by `epoll_wait()` call.
    pub fn wait(
        &self,
        max_events: usize,
        timeout: i32,
        events: &mut [EpollEvent],
    ) -> io::Result<usize> {
        // Safe because we give a valid epoll file descriptor and an array of epoll_event structures
        // that will be modified by the kernel to indicate information about the subset of file
        // descriptors in the interest list. We also check the return value.
        let events_count = SyscallReturnCode(unsafe {
            epoll_wait(
                self.epoll_fd,
                events.as_mut_ptr() as *mut epoll_event,
                max_events as i32,
                timeout,
            )
        })
        .into_result()? as usize;

        Ok(events_count)
    }
}

impl AsRawFd for Epoll {
    fn as_raw_fd(&self) -> RawFd {
        self.epoll_fd
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use utils::eventfd::EventFd;

    #[test]
    fn test_event_ops() {
        let mut event = EpollEvent::default();
        assert_eq!(event.events(), 0);
        assert_eq!(event.data(), 0);

        event = EpollEvent::new(EventSet::IN, 2);
        assert_eq!(event.events(), 1);
        assert_eq!(event.event_set(), EventSet::IN);

        assert_eq!(event.data(), 2);
        assert_eq!(event.fd(), 2);
    }

    #[test]
    fn test_epoll() {
        const DEFAULT__TIMEOUT: i32 = 250;
        const EVENT_BUFFER_SIZE: usize = 128;
        const MAX_EVENTS: usize = 10;

        let epoll = Epoll::new().unwrap();
        assert_eq!(epoll.epoll_fd, epoll.as_raw_fd());

        // Let's test different scenarios for `epoll_ctl()` and `epoll_wait()` functionality.

        let event_fd_1 = EventFd::new(libc::EFD_NONBLOCK).unwrap();
        // For EPOLLOUT to be available it is enough only to be possible to write a value of
        // at least 1 to the eventfd counter without blocking.
        // If we write a value greater than 0 to this counter, the fd will be available for
        // EPOLLIN events too.
        event_fd_1.write(1).unwrap();

        let mut event_1 =
            EpollEvent::new(EventSet::IN | EventSet::OUT, event_fd_1.as_raw_fd() as u64);

        // For EPOLL_CTL_ADD behavior we will try to add some fds with different event masks into
        // the interest list of epoll instance.
        assert!(epoll
            .ctl(
                ControlOperation::Add,
                event_fd_1.as_raw_fd() as i32,
                &event_1
            )
            .is_ok());

        // We can't add twice the same fd to epoll interest list.
        assert!(epoll
            .ctl(
                ControlOperation::Add,
                event_fd_1.as_raw_fd() as i32,
                &event_1
            )
            .is_err());

        let event_fd_2 = EventFd::new(libc::EFD_NONBLOCK).unwrap();
        event_fd_2.write(1).unwrap();
        assert!(epoll
            .ctl(
                ControlOperation::Add,
                event_fd_2.as_raw_fd() as i32,
                // For this fd, we want an Event instance that has `data` field set to other
                // value than the value of the fd and `events` without EPOLLIN type set.
                &EpollEvent::new(EventSet::OUT, 10)
            )
            .is_ok());

        // For the following eventfd we won't write anything to its counter, so we expect EPOLLIN
        // event to not be available for this fd, even if we say that we want to monitor this type
        // of event via EPOLL_CTL_ADD operation.
        let event_fd_3 = EventFd::new(libc::EFD_NONBLOCK).unwrap();
        let event_3 = EpollEvent::new(EventSet::OUT | EventSet::IN, event_fd_3.as_raw_fd() as u64);
        assert!(epoll
            .ctl(
                ControlOperation::Add,
                event_fd_3.as_raw_fd() as i32,
                &event_3
            )
            .is_ok());

        // Let's check `epoll_wait()` behavior for our epoll instance.
        let mut ready_events = vec![EpollEvent::default(); EVENT_BUFFER_SIZE];
        let mut ev_count = epoll
            .wait(MAX_EVENTS, DEFAULT__TIMEOUT, &mut ready_events[..])
            .unwrap();

        // We expect to have 3 fds in the ready list of epoll instance.
        assert_eq!(ev_count, 3);

        // Let's check also the Event values that are now returned in the ready list.
        assert_eq!(ready_events[0].data(), event_fd_1.as_raw_fd() as u64);
        // For this fd, `data` field was populated with random data instead of the
        // corresponding fd value.
        assert_eq!(ready_events[1].data(), 10);
        assert_eq!(ready_events[2].data(), event_fd_3.as_raw_fd() as u64);

        // EPOLLIN and EPOLLOUT should be available for this fd.
        assert_eq!(
            ready_events[0].events(),
            (EventSet::IN | EventSet::OUT).bits()
        );
        // Only EPOLLOUT is expected because we didn't want to monitor EPOLLIN on this fd.
        assert_eq!(ready_events[1].events(), EventSet::OUT.bits());
        // Only EPOLLOUT too because eventfd counter value is 0 (we didn't write a value
        // greater than 0 to it).
        assert_eq!(ready_events[2].events(), EventSet::OUT.bits());

        // Now we're gonna modify the Event instance for a fd to test EPOLL_CTL_MOD
        // behavior.
        // We create here a new Event with some events, other than those previously set,
        // that we want to monitor this time on event_fd_1.
        event_1 = EpollEvent::new(EventSet::OUT, 20);
        assert!(epoll
            .ctl(
                ControlOperation::Modify,
                event_fd_1.as_raw_fd() as i32,
                &event_1
            )
            .is_ok());

        let event_fd_4 = EventFd::new(libc::EFD_NONBLOCK).unwrap();
        // Can't modify a fd that wasn't added to epoll interest list.
        assert!(epoll
            .ctl(
                ControlOperation::Modify,
                event_fd_4.as_raw_fd() as i32,
                &EpollEvent::default()
            )
            .is_err());

        let _ = epoll
            .wait(MAX_EVENTS, DEFAULT__TIMEOUT, &mut ready_events[..])
            .unwrap();

        // Let's check that Event fields were indeed changed for the `event_fd_1` fd.
        assert_eq!(ready_events[0].data(), 20);
        // EPOLLOUT is now available for this fd as we've intended with EPOLL_CTL_MOD operation.
        assert_eq!(ready_events[0].events(), EventSet::OUT.bits());

        // Now let's set for a fd to not have any events monitored.
        assert!(epoll
            .ctl(
                ControlOperation::Modify,
                event_fd_1.as_raw_fd() as i32,
                &EpollEvent::default()
            )
            .is_ok());

        // In this particular case we expect to remain only with 2 fds in the ready list.
        ev_count = epoll
            .wait(MAX_EVENTS, DEFAULT__TIMEOUT, &mut ready_events[..])
            .unwrap();
        assert_eq!(ev_count, 2);

        // Let's also delete a fd from the interest list.
        assert!(epoll
            .ctl(
                ControlOperation::Delete,
                event_fd_2.as_raw_fd() as i32,
                &EpollEvent::default()
            )
            .is_ok());

        // We expect to have only one fd remained in the ready list (event_fd_3).
        ev_count = epoll
            .wait(MAX_EVENTS, DEFAULT__TIMEOUT, &mut ready_events[..])
            .unwrap();

        assert_eq!(ev_count, 1);
        assert_eq!(ready_events[0].data(), event_fd_3.as_raw_fd() as u64);
        assert_eq!(ready_events[0].events(), EventSet::OUT.bits());

        // If we try to remove a fd from epoll interest list that wasn't added before it will fail.
        assert!(epoll
            .ctl(
                ControlOperation::Delete,
                event_fd_4.as_raw_fd() as i32,
                &EpollEvent::default()
            )
            .is_err());
    }
}
<commit_msg>Do not implement Clone of Epoll<commit_after>// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use std::io;
use std::ops::Deref;
use std::os::unix::io::{AsRawFd, RawFd};

use libc::{
    epoll_create1, epoll_ctl, epoll_event, epoll_wait, EPOLLERR, EPOLLET, EPOLLEXCLUSIVE, EPOLLHUP,
    EPOLLIN, EPOLLONESHOT, EPOLLOUT, EPOLLPRI, EPOLLRDHUP, EPOLLWAKEUP, EPOLL_CLOEXEC,
    EPOLL_CTL_ADD, EPOLL_CTL_DEL, EPOLL_CTL_MOD,
};

use utils::syscall::SyscallReturnCode;

/// Wrapper over EPOLL_CTL_* operations that can be performed on a file descriptor.
#[repr(i32)]
pub enum ControlOperation {
    /// Add a file descriptor to the interest list.
    Add = EPOLL_CTL_ADD,
    /// Change the settings associated with a file descriptor that is
    /// already in the interest list.
    Modify = EPOLL_CTL_MOD,
    /// Remove a file descriptor from the interest list.
    Delete = EPOLL_CTL_DEL,
}

bitflags! {
    /// The type of events we can monitor a file descriptor for.
    pub struct EventSet: u32 {
        /// The associated file descriptor is available for read operations.
        const IN = EPOLLIN as u32;
        /// The associated file descriptor is available for write operations.
        const OUT = EPOLLOUT as u32;
        /// Error condition happened on the associated file descriptor.
        const ERROR = EPOLLERR as u32;
        /// This can be used to detect peer shutdown when using Edge Triggered monitoring.
        const READ_HANG_UP = EPOLLRDHUP as u32;
        /// Sets the Edge Triggered behavior for the associated file descriptor.
        /// The default behavior is Level Triggered.
        const EDGE_TRIGGERED = EPOLLET as u32;
        /// Hang up happened on the associated file descriptor. Note that `epoll_wait`
        /// will always wait for this event and it is not necessary to set it in events.
        const HANG_UP = EPOLLHUP as u32;
        /// There is an exceptional condition on that file descriptor. It is mostly used to
        /// set high priority for some data.
        const PRIORITY = EPOLLPRI as u32;
        /// The event is considered as being "processed" from the time when it is returned
        /// by a call to `epoll_wait` until the next call to `epoll_wait` on the same
        /// epoll file descriptor, the closure of that file descriptor, the removal of the
        /// event file descriptor via EPOLL_CTL_DEL, or the clearing of EPOLLWAKEUP
        /// for the event file descriptor via EPOLL_CTL_MOD.
        const WAKE_UP = EPOLLWAKEUP as u32;
        /// Sets the one-shot behavior for the associated file descriptor.
        const ONE_SHOT = EPOLLONESHOT as u32;
        /// Sets an exclusive wake up mode for the epoll file descriptor that is being
        /// attached to the associated file descriptor.
        /// When a wake up event occurs and multiple epoll file descriptors are attached to
        /// the same target file using this mode, one or more of the epoll file descriptors
        /// will receive an event with `epoll_wait`. The default here is for all those file
        /// descriptors to receive an event.
        const EXCLUSIVE = EPOLLEXCLUSIVE as u32;
    }
}

/// Wrapper over 'libc::epoll_event'.
///
// We are using `transparent` here to be super sure that this struct and its fields
// have the same alignment as those from the `epoll_event` struct from C.
#[repr(transparent)]
#[derive(Clone)]
pub struct EpollEvent(epoll_event);

impl Deref for EpollEvent {
    type Target = epoll_event;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl Default for EpollEvent {
    fn default() -> Self {
        EpollEvent(epoll_event {
            events: 0u32,
            u64: 0u64,
        })
    }
}

impl EpollEvent {
    /// Create a new epoll_event instance with the following fields: `events`, which contains
    /// an event mask and `data` which represents a user data variable. `data` field can be
    /// a fd on which we want to monitor the events specified by `events`.
    pub fn new(events: EventSet, data: u64) -> Self {
        EpollEvent(epoll_event {
            events: events.bits(),
            u64: data,
        })
    }

    /// Returns the `events` from `libc::epoll_event`.
    pub fn events(&self) -> u32 {
        self.events
    }

    /// Returns the `EventSet` corresponding to `epoll_event.events`.
    ///
    /// # Panics
    ///
    /// Panics if `libc::epoll_event` contains invalid events.
    pub fn event_set(&self) -> EventSet {
        // This unwrap is safe because `epoll_events` can only be user created or
        // initialized by the kernel.We trust the kernel to only send us valid
        // events. The user can only initialize `epoll_events` using valid events.
        EventSet::from_bits(self.events()).unwrap()
    }

    /// Returns the `data` from the `libc::epoll_event`.
    pub fn data(&self) -> u64 {
        self.u64
    }

    /// Converts the `libc::epoll_event` data to a RawFd.
    ///
    /// This conversion is lossy when the data does not correspond to a RawFd
    /// (data does not fit in a i32).
    pub fn fd(&self) -> RawFd {
        self.u64 as i32
    }
}

/// Wrapper over epoll functionality.
#[derive(Debug)]
pub struct Epoll {
    epoll_fd: RawFd,
}

impl Epoll {
    /// Create a new epoll file descriptor.
    pub fn new() -> io::Result<Self> {
        let epoll_fd = SyscallReturnCode(unsafe { epoll_create1(EPOLL_CLOEXEC) }).into_result()?;
        Ok(Epoll { epoll_fd })
    }

    /// Wrapper for `libc::epoll_ctl`.
    ///
    /// This can be used for adding, modifying or removing a file descriptor in the
    /// interest list of the epoll instance.
    ///
    /// # Arguments
    ///
    /// * `operation` refers to the action to be performed on the file descriptor.
    /// * `fd` is the file descriptor on which we want to perform `operation`.
    /// * `event` refers to the `epoll_event` instance that is linked to `fd`.
    pub fn ctl(
        &self,
        operation: ControlOperation,
        fd: RawFd,
        event: &EpollEvent,
    ) -> io::Result<()> {
        // Safe because we give a valid epoll file descriptor, a valid file descriptor to watch,
        // as well as a valid epoll_event structure. We also check the return value.
        SyscallReturnCode(unsafe {
            epoll_ctl(
                self.epoll_fd,
                operation as i32,
                fd,
                event as *const EpollEvent as *mut epoll_event,
            )
        })
        .into_empty_result()
    }

    /// Wrapper for `libc::epoll_wait`.
    /// Returns the number of file descriptors in the interest list that became ready
    /// for I/O or `errno` if an error occurred.
    ///
    /// # Arguments
    ///
    /// * `max_events` is the maximum number of events that we want to be returned in
    /// `events` buffer.
    /// * `timeout` specifies for how long the `epoll_wait` system call will block
    /// (measured in milliseconds).
    /// * `events` points to a memory area that will be used for storing the events
    /// returned by `epoll_wait()` call.
    pub fn wait(
        &self,
        max_events: usize,
        timeout: i32,
        events: &mut [EpollEvent],
    ) -> io::Result<usize> {
        // Safe because we give a valid epoll file descriptor and an array of epoll_event structures
        // that will be modified by the kernel to indicate information about the subset of file
        // descriptors in the interest list. We also check the return value.
        let events_count = SyscallReturnCode(unsafe {
            epoll_wait(
                self.epoll_fd,
                events.as_mut_ptr() as *mut epoll_event,
                max_events as i32,
                timeout,
            )
        })
        .into_result()? as usize;

        Ok(events_count)
    }
}

impl AsRawFd for Epoll {
    fn as_raw_fd(&self) -> RawFd {
        self.epoll_fd
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use utils::eventfd::EventFd;

    #[test]
    fn test_event_ops() {
        let mut event = EpollEvent::default();
        assert_eq!(event.events(), 0);
        assert_eq!(event.data(), 0);

        event = EpollEvent::new(EventSet::IN, 2);
        assert_eq!(event.events(), 1);
        assert_eq!(event.event_set(), EventSet::IN);

        assert_eq!(event.data(), 2);
        assert_eq!(event.fd(), 2);
    }

    #[test]
    fn test_epoll() {
        const DEFAULT__TIMEOUT: i32 = 250;
        const EVENT_BUFFER_SIZE: usize = 128;
        const MAX_EVENTS: usize = 10;

        let epoll = Epoll::new().unwrap();
        assert_eq!(epoll.epoll_fd, epoll.as_raw_fd());

        // Let's test different scenarios for `epoll_ctl()` and `epoll_wait()` functionality.

        let event_fd_1 = EventFd::new(libc::EFD_NONBLOCK).unwrap();
        // For EPOLLOUT to be available it is enough only to be possible to write a value of
        // at least 1 to the eventfd counter without blocking.
        // If we write a value greater than 0 to this counter, the fd will be available for
        // EPOLLIN events too.
        event_fd_1.write(1).unwrap();

        let mut event_1 =
            EpollEvent::new(EventSet::IN | EventSet::OUT, event_fd_1.as_raw_fd() as u64);

        // For EPOLL_CTL_ADD behavior we will try to add some fds with different event masks into
        // the interest list of epoll instance.
        assert!(epoll
            .ctl(
                ControlOperation::Add,
                event_fd_1.as_raw_fd() as i32,
                &event_1
            )
            .is_ok());

        // We can't add twice the same fd to epoll interest list.
        assert!(epoll
            .ctl(
                ControlOperation::Add,
                event_fd_1.as_raw_fd() as i32,
                &event_1
            )
            .is_err());

        let event_fd_2 = EventFd::new(libc::EFD_NONBLOCK).unwrap();
        event_fd_2.write(1).unwrap();
        assert!(epoll
            .ctl(
                ControlOperation::Add,
                event_fd_2.as_raw_fd() as i32,
                // For this fd, we want an Event instance that has `data` field set to other
                // value than the value of the fd and `events` without EPOLLIN type set.
                &EpollEvent::new(EventSet::OUT, 10)
            )
            .is_ok());

        // For the following eventfd we won't write anything to its counter, so we expect EPOLLIN
        // event to not be available for this fd, even if we say that we want to monitor this type
        // of event via EPOLL_CTL_ADD operation.
        let event_fd_3 = EventFd::new(libc::EFD_NONBLOCK).unwrap();
        let event_3 = EpollEvent::new(EventSet::OUT | EventSet::IN, event_fd_3.as_raw_fd() as u64);
        assert!(epoll
            .ctl(
                ControlOperation::Add,
                event_fd_3.as_raw_fd() as i32,
                &event_3
            )
            .is_ok());

        // Let's check `epoll_wait()` behavior for our epoll instance.
        let mut ready_events = vec![EpollEvent::default(); EVENT_BUFFER_SIZE];
        let mut ev_count = epoll
            .wait(MAX_EVENTS, DEFAULT__TIMEOUT, &mut ready_events[..])
            .unwrap();

        // We expect to have 3 fds in the ready list of epoll instance.
        assert_eq!(ev_count, 3);

        // Let's check also the Event values that are now returned in the ready list.
        assert_eq!(ready_events[0].data(), event_fd_1.as_raw_fd() as u64);
        // For this fd, `data` field was populated with random data instead of the
        // corresponding fd value.
        assert_eq!(ready_events[1].data(), 10);
        assert_eq!(ready_events[2].data(), event_fd_3.as_raw_fd() as u64);

        // EPOLLIN and EPOLLOUT should be available for this fd.
        assert_eq!(
            ready_events[0].events(),
            (EventSet::IN | EventSet::OUT).bits()
        );
        // Only EPOLLOUT is expected because we didn't want to monitor EPOLLIN on this fd.
        assert_eq!(ready_events[1].events(), EventSet::OUT.bits());
        // Only EPOLLOUT too because eventfd counter value is 0 (we didn't write a value
        // greater than 0 to it).
        assert_eq!(ready_events[2].events(), EventSet::OUT.bits());

        // Now we're gonna modify the Event instance for a fd to test EPOLL_CTL_MOD
        // behavior.
        // We create here a new Event with some events, other than those previously set,
        // that we want to monitor this time on event_fd_1.
        event_1 = EpollEvent::new(EventSet::OUT, 20);
        assert!(epoll
            .ctl(
                ControlOperation::Modify,
                event_fd_1.as_raw_fd() as i32,
                &event_1
            )
            .is_ok());

        let event_fd_4 = EventFd::new(libc::EFD_NONBLOCK).unwrap();
        // Can't modify a fd that wasn't added to epoll interest list.
        assert!(epoll
            .ctl(
                ControlOperation::Modify,
                event_fd_4.as_raw_fd() as i32,
                &EpollEvent::default()
            )
            .is_err());

        let _ = epoll
            .wait(MAX_EVENTS, DEFAULT__TIMEOUT, &mut ready_events[..])
            .unwrap();

        // Let's check that Event fields were indeed changed for the `event_fd_1` fd.
        assert_eq!(ready_events[0].data(), 20);
        // EPOLLOUT is now available for this fd as we've intended with EPOLL_CTL_MOD operation.
        assert_eq!(ready_events[0].events(), EventSet::OUT.bits());

        // Now let's set for a fd to not have any events monitored.
        assert!(epoll
            .ctl(
                ControlOperation::Modify,
                event_fd_1.as_raw_fd() as i32,
                &EpollEvent::default()
            )
            .is_ok());

        // In this particular case we expect to remain only with 2 fds in the ready list.
        ev_count = epoll
            .wait(MAX_EVENTS, DEFAULT__TIMEOUT, &mut ready_events[..])
            .unwrap();
        assert_eq!(ev_count, 2);

        // Let's also delete a fd from the interest list.
        assert!(epoll
            .ctl(
                ControlOperation::Delete,
                event_fd_2.as_raw_fd() as i32,
                &EpollEvent::default()
            )
            .is_ok());

        // We expect to have only one fd remained in the ready list (event_fd_3).
        ev_count = epoll
            .wait(MAX_EVENTS, DEFAULT__TIMEOUT, &mut ready_events[..])
            .unwrap();

        assert_eq!(ev_count, 1);
        assert_eq!(ready_events[0].data(), event_fd_3.as_raw_fd() as u64);
        assert_eq!(ready_events[0].events(), EventSet::OUT.bits());

        // If we try to remove a fd from epoll interest list that wasn't added before it will fail.
        assert!(epoll
            .ctl(
                ControlOperation::Delete,
                event_fd_4.as_raw_fd() as i32,
                &EpollEvent::default()
            )
            .is_err());
    }
}
<|endoftext|>