<commit_before>#![feature(alloc_system)]
extern crate alloc_system;

extern crate env_logger;
extern crate futures;
extern crate rand;
extern crate rdkafka;

mod test_utils;

use futures::*;

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

use test_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 mut consumer = create_stream_consumer(&rand_test_group(), None);
    consumer.subscribe(&vec![topic_name.as_str()]).unwrap();

    let _consumer_future = consumer.start()
        .take(100)
        .for_each(|message| {
            match message {
                Ok(m) => {
                    let id = message_map.get(&(m.partition(), m.offset())).unwrap();
                    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_name(), 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 mut 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 mut consumer = create_stream_consumer(&rand_test_group(), None);
    consumer.subscribe(&vec![topic_name.as_str()]).unwrap();

    let _consumer_future = consumer.start()
        .take(100)
        .for_each(|message| {
            match message {
                Ok(m) => {
                    let id = message_map.get(&(m.partition(), m.offset())).unwrap();
                    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, 100).unwrap();
    let tp = tpl.find_partition(&topic_name, 0).unwrap();
    assert_eq!(tp.offset(), Offset::Offset(100));
}

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

    let mut 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 == 5 {
                    break;
                }
            }
            Ok(m) => panic!("A message was actually received: {:?}", m),
            Err(e) => panic!("Unexpected error while receiving message: {:?}", e)
        };
    }

    assert_eq!(timeouts_count, 5);
    // It should take 800ms
    assert!(Instant::now().duration_since(first_poll_time.unwrap()) < Duration::from_millis(1000));
    assert!(Instant::now().duration_since(first_poll_time.unwrap()) > Duration::from_millis(600));
}

// METADATA

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

    let topic_name = rand_test_topic();
    produce_messages(&topic_name, 1, &value_fn, &key_fn, Some(0), None);
    produce_messages(&topic_name, 1, &value_fn, &key_fn, Some(1), None);
    produce_messages(&topic_name, 1, &value_fn, &key_fn, Some(2), None);
    let consumer = create_stream_consumer(&rand_test_group(), None);

    let metadata = consumer.fetch_metadata(None, 5000).unwrap();

    let topic_metadata = metadata.topics().iter()
        .find(|m| m.name() == topic_name).unwrap();

    let mut ids = topic_metadata.partitions().iter().map(|p| p.id()).collect::<Vec<_>>();
    ids.sort();

    assert_eq!(ids, vec![0, 1, 2]);
    // assert_eq!(topic_metadata.error(), None);
    assert_eq!(topic_metadata.partitions().len(), 3);
    assert_eq!(topic_metadata.partitions()[0].leader(), 0);
    assert_eq!(topic_metadata.partitions()[1].leader(), 0);
    assert_eq!(topic_metadata.partitions()[2].leader(), 0);
    assert_eq!(topic_metadata.partitions()[0].replicas(), &[0]);
    assert_eq!(topic_metadata.partitions()[0].isr(), &[0]);

    let metadata_one_topic = consumer.fetch_metadata(Some(&topic_name), 5000).unwrap();
    assert_eq!(metadata_one_topic.topics().len(), 1);
}

// 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 mut consumer = create_stream_consumer(&rand_test_group(), None);
    consumer.subscribe(&vec![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 mut consumer = create_stream_consumer(&rand_test_group(), Some(config));
    consumer.subscribe(&vec![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(None, 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());
}

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

    let topic_name = rand_test_topic();
    produce_messages(&topic_name, 10, &value_fn, &key_fn, None, None);
    let mut consumer = create_stream_consumer(&rand_test_group(), None);
    consumer.subscribe(&vec![topic_name.as_str()]).unwrap();

    let _consumer_future = consumer.start().take(10).wait();

    let mut tpl = TopicPartitionList::new();
    tpl.add_topic_unassigned(&topic_name);
    assert_eq!(tpl, consumer.subscription().unwrap());
}

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

    let topic_name = rand_test_topic();
    let group_name = rand_test_group();
    produce_messages(&topic_name, 1, &value_fn, &key_fn, Some(0), None);
    produce_messages(&topic_name, 1, &value_fn, &key_fn, Some(1), None);
    produce_messages(&topic_name, 1, &value_fn, &key_fn, Some(2), None);
    let mut consumer = create_stream_consumer(&group_name, None);
    consumer.subscribe(&vec![topic_name.as_str()]).unwrap();

    // Make sure the consumer joins the group
    let _consumer_future = consumer.start()
        .take(1)
        .for_each(|_| Ok(()))
        .wait();

    let group_list = consumer.fetch_group_list(None, 5000).unwrap();

    // Print all the data, valgrind will check memory access
    for group in group_list.groups().iter() {
        println!("{} {} {} {}", group.name(), group.state(), group.protocol(), group.protocol_type());
        for member in group.members() {
            println!("  {} {} {}", member.id(), member.client_id(), member.client_host());
        }
    }

    let group_list2 = consumer.fetch_group_list(Some(&group_name), 5000).unwrap();
    assert_eq!(group_list2.groups().len(), 1);

    let consumer_group = group_list2.groups().iter().find(|&g| g.name() == group_name).unwrap();
    assert_eq!(consumer_group.members().len(), 1);

    let consumer_member = &consumer_group.members()[0];
    assert_eq!(consumer_member.client_id(), "rdkafka_integration_test_client");
}
<commit_msg>Relax time constraint to fix flaky test<commit_after>#![feature(alloc_system)]
extern crate alloc_system;

extern crate env_logger;
extern crate futures;
extern crate rand;
extern crate rdkafka;

mod test_utils;

use futures::*;

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

use test_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 mut consumer = create_stream_consumer(&rand_test_group(), None);
    consumer.subscribe(&vec![topic_name.as_str()]).unwrap();

    let _consumer_future = consumer.start()
        .take(100)
        .for_each(|message| {
            match message {
                Ok(m) => {
                    let id = message_map.get(&(m.partition(), m.offset())).unwrap();
                    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_name(), 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 mut 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 mut consumer = create_stream_consumer(&rand_test_group(), None);
    consumer.subscribe(&vec![topic_name.as_str()]).unwrap();

    let _consumer_future = consumer.start()
        .take(100)
        .for_each(|message| {
            match message {
                Ok(m) => {
                    let id = message_map.get(&(m.partition(), m.offset())).unwrap();
                    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, 100).unwrap();
    let tp = tpl.find_partition(&topic_name, 0).unwrap();
    assert_eq!(tp.offset(), Offset::Offset(100));
}

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

    let mut 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 == 5 {
                    break;
                }
            }
            Ok(m) => panic!("A message was actually received: {:?}", m),
            Err(e) => panic!("Unexpected error while receiving message: {:?}", e)
        };
    }

    assert_eq!(timeouts_count, 5);
    // It should take 800ms
    assert!(Instant::now().duration_since(first_poll_time.unwrap()) < Duration::from_millis(1500));
    assert!(Instant::now().duration_since(first_poll_time.unwrap()) > Duration::from_millis(600));
}

// METADATA

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

    let topic_name = rand_test_topic();
    produce_messages(&topic_name, 1, &value_fn, &key_fn, Some(0), None);
    produce_messages(&topic_name, 1, &value_fn, &key_fn, Some(1), None);
    produce_messages(&topic_name, 1, &value_fn, &key_fn, Some(2), None);
    let consumer = create_stream_consumer(&rand_test_group(), None);

    let metadata = consumer.fetch_metadata(None, 5000).unwrap();

    let topic_metadata = metadata.topics().iter()
        .find(|m| m.name() == topic_name).unwrap();

    let mut ids = topic_metadata.partitions().iter().map(|p| p.id()).collect::<Vec<_>>();
    ids.sort();

    assert_eq!(ids, vec![0, 1, 2]);
    // assert_eq!(topic_metadata.error(), None);
    assert_eq!(topic_metadata.partitions().len(), 3);
    assert_eq!(topic_metadata.partitions()[0].leader(), 0);
    assert_eq!(topic_metadata.partitions()[1].leader(), 0);
    assert_eq!(topic_metadata.partitions()[2].leader(), 0);
    assert_eq!(topic_metadata.partitions()[0].replicas(), &[0]);
    assert_eq!(topic_metadata.partitions()[0].isr(), &[0]);

    let metadata_one_topic = consumer.fetch_metadata(Some(&topic_name), 5000).unwrap();
    assert_eq!(metadata_one_topic.topics().len(), 1);
}

// 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 mut consumer = create_stream_consumer(&rand_test_group(), None);
    consumer.subscribe(&vec![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 mut consumer = create_stream_consumer(&rand_test_group(), Some(config));
    consumer.subscribe(&vec![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(None, 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());
}

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

    let topic_name = rand_test_topic();
    produce_messages(&topic_name, 10, &value_fn, &key_fn, None, None);
    let mut consumer = create_stream_consumer(&rand_test_group(), None);
    consumer.subscribe(&vec![topic_name.as_str()]).unwrap();

    let _consumer_future = consumer.start().take(10).wait();

    let mut tpl = TopicPartitionList::new();
    tpl.add_topic_unassigned(&topic_name);
    assert_eq!(tpl, consumer.subscription().unwrap());
}

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

    let topic_name = rand_test_topic();
    let group_name = rand_test_group();
    produce_messages(&topic_name, 1, &value_fn, &key_fn, Some(0), None);
    produce_messages(&topic_name, 1, &value_fn, &key_fn, Some(1), None);
    produce_messages(&topic_name, 1, &value_fn, &key_fn, Some(2), None);
    let mut consumer = create_stream_consumer(&group_name, None);
    consumer.subscribe(&vec![topic_name.as_str()]).unwrap();

    // Make sure the consumer joins the group
    let _consumer_future = consumer.start()
        .take(1)
        .for_each(|_| Ok(()))
        .wait();

    let group_list = consumer.fetch_group_list(None, 5000).unwrap();

    // Print all the data, valgrind will check memory access
    for group in group_list.groups().iter() {
        println!("{} {} {} {}", group.name(), group.state(), group.protocol(), group.protocol_type());
        for member in group.members() {
            println!("  {} {} {}", member.id(), member.client_id(), member.client_host());
        }
    }

    let group_list2 = consumer.fetch_group_list(Some(&group_name), 5000).unwrap();
    assert_eq!(group_list2.groups().len(), 1);

    let consumer_group = group_list2.groups().iter().find(|&g| g.name() == group_name).unwrap();
    assert_eq!(consumer_group.members().len(), 1);

    let consumer_member = &consumer_group.members()[0];
    assert_eq!(consumer_member.client_id(), "rdkafka_integration_test_client");
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate serde_derive;

#[macro_use]
extern crate structopt;
extern crate tabin_plugins;

use structopt::StructOpt;

use tabin_plugins::linux::pages_to_human_size;
use tabin_plugins::Status;
use tabin_plugins::procfs::{LoadProcsError, MemInfo, ProcFsError, RunningProcs};

/// Check the ram usage of the current computer
#[derive(Deserialize, StructOpt, Debug)]
#[structopt(name = "check-ram (part of tabin-plugins)",
            raw(setting = "structopt::clap::AppSettings::ColoredHelp"))]
struct Args {
    #[structopt(short = "w", long = "warn", help = "Percent to warn at", default_value = "85")]
    warn: f64,
    #[structopt(short = "c", long = "crit", help = "Percent to go critical at",
                default_value = "95")]
    crit: f64,

    #[structopt(long = "show-hogs", name = "count",
                help = "Show <count> most ram-intensive processes in this computer.",
                default_value = "0")]
    show_hogs: usize,
}

#[cfg_attr(test, allow(dead_code))]
fn main() {
    let args = Args::from_args();
    let mem = MemInfo::load();
    let status = compare_status(args.crit, args.warn, &mem);
    if args.show_hogs > 0 {
        let mut load_errors = None;
        let per_proc = match RunningProcs::currently_running() {
            Ok(procs) => procs,
            Err(ProcFsError::LoadProcsError(LoadProcsError { procs, errors })) => {
                load_errors = Some(errors);
                procs
            }
            Err(err) => {
                eprintln!("UNKNOWN: Unexpected error loading procs: {}", err);
                RunningProcs::empty()
            }
        };

        let mut procs = per_proc.0.values().collect::<Vec<_>>();
        procs.sort_by(|l, r| r.stat.rss.cmp(&l.stat.rss));
        println!(
            "INFO [check-ram]: {} processes running, top {} ram hogs:",
            procs.len(), args.show_hogs
        );
        for process in procs.iter().take(args.show_hogs) {
            let system_kb = mem.total.unwrap();
            let percent = process.percent_ram(system_kb * 1024);
            println!(
                "[{:>6}]{:>5.1}% {:>6}: {}",
                process.stat.pid,
                percent,
                pages_to_human_size(process.stat.rss),
                process.useful_cmdline()
            );
        }
        if let Some(errors) = load_errors {
            for err in errors {
                eprintln!("UNKNOWN: error loading process: {}", err);
            }
        }
    };
    status.exit();
}

fn compare_status(crit: f64, warn: f64, mem: &MemInfo) -> Status {
    match mem.percent_used() {
        Ok(percent) => if percent > crit {
            println!("CRITICAL [check-ram]: {:.1}% > {}%", percent, crit);
            Status::Critical
        } else if percent > warn {
            println!("WARNING [check-ram]: {:.1}% > {}%", percent, warn);
            Status::Warning
        } else {
            println!("OK [check-ram]: {:.1}% < {}%", percent, warn);
            Status::Ok
        },
        Err(e) => {
            println!("UNKNOWN [check-ram]: UNEXPECTED ERROR {:?}", e);
            Status::Unknown
        }
    }
}

#[cfg(test)]
mod test {
    use super::compare_status;
    use tabin_plugins::Status;
    use tabin_plugins::procfs::MemInfo;

    #[test]
    fn alerts_when_told_to() {
        let mem = MemInfo {
            total: Some(100),
            available: Some(15), // 15% free means 85% used
            free: None,
            cached: None,
        };
        let crit_threshold = 80.0;

        assert_eq!(compare_status(crit_threshold, 25.0, &mem), Status::Critical);
    }
}
<commit_msg>check-ram: rustfmt<commit_after>#[macro_use]
extern crate serde_derive;

#[macro_use]
extern crate structopt;
extern crate tabin_plugins;

use structopt::StructOpt;

use tabin_plugins::linux::pages_to_human_size;
use tabin_plugins::Status;
use tabin_plugins::procfs::{LoadProcsError, MemInfo, ProcFsError, RunningProcs};

/// Check the ram usage of the current computer
#[derive(Deserialize, StructOpt, Debug)]
#[structopt(name = "check-ram (part of tabin-plugins)",
            raw(setting = "structopt::clap::AppSettings::ColoredHelp"))]
struct Args {
    #[structopt(short = "w", long = "warn", help = "Percent to warn at", default_value = "85")]
    warn: f64,
    #[structopt(short = "c", long = "crit", help = "Percent to go critical at",
                default_value = "95")]
    crit: f64,

    #[structopt(long = "show-hogs", name = "count",
                help = "Show <count> most ram-intensive processes in this computer.",
                default_value = "0")]
    show_hogs: usize,
}

#[cfg_attr(test, allow(dead_code))]
fn main() {
    let args = Args::from_args();
    let mem = MemInfo::load();
    let status = compare_status(args.crit, args.warn, &mem);
    if args.show_hogs > 0 {
        let mut load_errors = None;
        let per_proc = match RunningProcs::currently_running() {
            Ok(procs) => procs,
            Err(ProcFsError::LoadProcsError(LoadProcsError { procs, errors })) => {
                load_errors = Some(errors);
                procs
            }
            Err(err) => {
                eprintln!("UNKNOWN: Unexpected error loading procs: {}", err);
                RunningProcs::empty()
            }
        };

        let mut procs = per_proc.0.values().collect::<Vec<_>>();
        procs.sort_by(|l, r| r.stat.rss.cmp(&l.stat.rss));
        println!(
            "INFO [check-ram]: {} processes running, top {} ram hogs:",
            procs.len(),
            args.show_hogs
        );
        for process in procs.iter().take(args.show_hogs) {
            let system_kb = mem.total.unwrap();
            let percent = process.percent_ram(system_kb * 1024);
            println!(
                "[{:>6}]{:>5.1}% {:>6}: {}",
                process.stat.pid,
                percent,
                pages_to_human_size(process.stat.rss),
                process.useful_cmdline()
            );
        }
        if let Some(errors) = load_errors {
            for err in errors {
                eprintln!("UNKNOWN: error loading process: {}", err);
            }
        }
    };
    status.exit();
}

fn compare_status(crit: f64, warn: f64, mem: &MemInfo) -> Status {
    match mem.percent_used() {
        Ok(percent) => if percent > crit {
            println!("CRITICAL [check-ram]: {:.1}% > {}%", percent, crit);
            Status::Critical
        } else if percent > warn {
            println!("WARNING [check-ram]: {:.1}% > {}%", percent, warn);
            Status::Warning
        } else {
            println!("OK [check-ram]: {:.1}% < {}%", percent, warn);
            Status::Ok
        },
        Err(e) => {
            println!("UNKNOWN [check-ram]: UNEXPECTED ERROR {:?}", e);
            Status::Unknown
        }
    }
}

#[cfg(test)]
mod test {
    use super::compare_status;
    use tabin_plugins::Status;
    use tabin_plugins::procfs::MemInfo;

    #[test]
    fn alerts_when_told_to() {
        let mem = MemInfo {
            total: Some(100),
            available: Some(15), // 15% free means 85% used
            free: None,
            cached: None,
        };
        let crit_threshold = 80.0;

        assert_eq!(compare_status(crit_threshold, 25.0, &mem), Status::Critical);
    }
}
<|endoftext|>


<commit_before>use std::collections::BTreeMap;
use std::sync::{Arc, RwLock};

use exonum::storage::{Database, Patch, Change, Result as StorageResult, Snapshot};

/// Implementation of a `Database`, which allows to rollback commits introduced by the `merge()`
/// function.
///
/// **Note:** Intended for testing purposes only. Probably inefficient.
#[derive(Clone, Debug)]
pub struct CheckpointDb<T> {
    inner: T,
    journal: Arc<RwLock<Vec<Patch>>>,
}

impl<T: Database + Clone> CheckpointDb<T> {
    /// Creates a new checkpointed database that uses the specified `db` as the underlying
    /// data storage.
    pub fn new(db: T) -> Self {
        CheckpointDb {
            inner: db,
            journal: Arc::new(RwLock::new(Vec::new())),
        }
    }

    /// Rolls this database by udoing the latest `count` `merge()` operations.
    pub fn rollback(&mut self, count: usize) -> bool {
        let journal_len = self.journal
            .read()
            .expect("Cannot acquire read lock on journal")
            .len();
        assert!(
            journal_len >= count,
            "Cannot rollback {} changes; only {} checkpoints in the journal",
            count,
            journal_len
        );

        let mut journal = self.journal.write().expect(
            "Cannot acquire write lock on journal",
        );

        for _ in 0..count {
            if let Some(patch) = journal.pop() {
                self.inner.merge(patch).expect(
                    "Cannot merge roll-back patch",
                );
            } else {
                panic!("Cannot rollback the database, the journal is empty");
            }
        }

        journal_len < count
    }

    /// Returns a handler to the database. The handler could be used to roll the database back
    /// without having the ownership to it.
    pub fn handler(&self) -> CheckpointDbHandler<T> {
        CheckpointDbHandler(Clone::clone(self))
    }
}

impl<T: Database + Clone> Database for CheckpointDb<T> {
    fn clone(&self) -> Box<Database> {
        Box::new(Clone::clone(self))
    }

    fn snapshot(&self) -> Box<Snapshot> {
        self.inner.snapshot()
    }

    fn merge(&mut self, patch: Patch) -> StorageResult<()> {
        let snapshot = self.inner.snapshot();
        let mut rev_patch = Patch::new();

        for (name, kv_map) in &patch {
            let mut rev_kv_map = BTreeMap::new();
            for key in kv_map.keys() {
                match snapshot.get(name, key) {
                    Some(value) => {
                        rev_kv_map.insert(key.clone(), Change::Put(value));
                    }
                    None => {
                        rev_kv_map.insert(key.clone(), Change::Delete);
                    }
                }
            }
            rev_patch.insert(name.to_string(), rev_kv_map);
        }

        {
            let mut journal = self.journal.write().expect(
                "Cannot acquire write lock on journal",
            );
            journal.push(rev_patch);
        }
        self.inner.merge(patch)
    }

    fn merge_sync(&mut self, patch: Patch) -> StorageResult<()> {
        self.merge(patch)
    }
}

/// Handler to a checkpointed database.
#[derive(Clone, Debug)]
pub struct CheckpointDbHandler<T>(CheckpointDb<T>);

impl<T: Database + Clone> CheckpointDbHandler<T> {
    /// Rolls this database by udoing the latest `count` `merge()` operations.
    pub fn rollback(&mut self, count: usize) -> bool {
        self.0.rollback(count)
    }
}

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

    fn create_patch<'a, I>(from: I) -> Patch
    where
        I: IntoIterator<Item = (&'a str, Vec<u8>, Change)>,
    {
        let mut patch = Patch::new();

        for (name, key, value) in from.into_iter() {
            let name = name.to_string();
            let kv_map = patch.entry(name).or_insert_with(BTreeMap::new);
            if kv_map.insert(key, value).is_some() {
                panic!("Attempt to re-insert a key into patch");
            }
        }

        patch
    }

    #[test]
    fn test_checkpointdb_basics() {
        let mut db = CheckpointDb::new(MemoryDB::new());
        let mut fork = db.fork();
        fork.put("foo", vec![], vec![2]);
        db.merge(fork.into_patch()).unwrap();
        {
            let journal = db.journal.read().expect(
                "Cannot acquire read lock on journal",
            );
            assert_eq!(journal.len(), 1);
            assert_eq!(
                journal[0],
                create_patch(vec![("foo", vec![], Change::Delete)])
            );
        }

        let snapshot = db.snapshot();
        assert_eq!(snapshot.get("foo", &[]), Some(vec![2]));

        let mut fork = db.fork();
        fork.put("foo", vec![], vec![3]);
        fork.put("bar", vec![1], vec![4]);
        db.merge(fork.into_patch()).unwrap();
        {
            let journal = db.journal.read().expect(
                "Cannot acquire read lock on journal",
            );
            assert_eq!(journal.len(), 2);
            assert_eq!(
                journal[0],
                create_patch(vec![("foo", vec![], Change::Delete)])
            );
            assert_eq!(
                journal[1],
                create_patch(vec![
                    ("foo", vec![], Change::Put(vec![2])),
                    ("bar", vec![1], Change::Delete),
                ])
            );
        }

        // Check that the old snapshot still corresponds to the same DB state
        assert_eq!(snapshot.get("foo", &[]), Some(vec![2]));
        let snapshot = db.snapshot();
        assert_eq!(snapshot.get("foo", &[]), Some(vec![3]));
    }

    #[test]
    fn test_checkpointdb_rollback() {
        let mut db = CheckpointDb::new(MemoryDB::new());
        let mut fork = db.fork();
        fork.put("foo", vec![], vec![2]);
        db.merge(fork.into_patch()).unwrap();

        let mut fork = db.fork();
        fork.put("foo", vec![], vec![3]);
        fork.put("bar", vec![1], vec![4]);
        db.merge(fork.into_patch()).unwrap();
        {
            let journal = db.journal.read().expect(
                "Cannot acquire read lock on journal",
            );
            assert_eq!(journal.len(), 2);
        }

        let snapshot = db.snapshot();
        assert_eq!(snapshot.get("foo", &[]), Some(vec![3]));
        db.rollback(1);
        let snapshot = db.snapshot();
        assert_eq!(snapshot.get("foo", &[]), Some(vec![2]));
        assert_eq!(snapshot.get("bar", &[1]), None);
        {
            let journal = db.journal.read().expect(
                "Cannot acquire read lock on journal",
            );
            assert_eq!(journal.len(), 1);
        }

        // Check that DB continues working as usual after a rollback
        let mut fork = db.fork();
        fork.put("foo", vec![], vec![4]);
        fork.put("foo", vec![0, 0], vec![255]);
        db.merge(fork.into_patch()).unwrap();
        {
            let journal = db.journal.read().expect(
                "Cannot acquire read lock on journal",
            );
            assert_eq!(journal.len(), 2);
        }
        let snapshot = db.snapshot();
        assert_eq!(snapshot.get("foo", &[]), Some(vec![4]));
        assert_eq!(snapshot.get("foo", &[0, 0]), Some(vec![255]));

        let mut fork = db.fork();
        fork.put("bar", vec![1], vec![254]);
        db.merge(fork.into_patch()).unwrap();
        {
            let journal = db.journal.read().expect(
                "Cannot acquire read lock on journal",
            );
            assert_eq!(journal.len(), 3);
        }
        let new_snapshot = db.snapshot();
        assert_eq!(new_snapshot.get("foo", &[]), Some(vec![4]));
        assert_eq!(new_snapshot.get("foo", &[0, 0]), Some(vec![255]));
        assert_eq!(new_snapshot.get("bar", &[1]), Some(vec![254]));

        db.rollback(2);
        {
            let journal = db.journal.read().expect(
                "Cannot acquire read lock on journal",
            );
            assert_eq!(journal.len(), 1);
        }
        let snapshot = db.snapshot();
        assert_eq!(snapshot.get("foo", &[]), Some(vec![2]));
        assert_eq!(snapshot.get("foo", &[0, 0]), None);
        assert_eq!(snapshot.get("bar", &[1]), None);

        assert_eq!(new_snapshot.get("foo", &[]), Some(vec![4]));
        assert_eq!(new_snapshot.get("foo", &[0, 0]), Some(vec![255]));
        assert_eq!(new_snapshot.get("bar", &[1]), Some(vec![254]));
    }

    #[test]
    fn test_checkpointdb_clone() {
        let mut db = CheckpointDb::new(MemoryDB::new());
        let mut db_clone = Database::clone(&db);

        let mut fork = db.fork();
        fork.put("foo", vec![], vec![2]);
        db.merge(fork.into_patch()).unwrap();

        let snapshot = db_clone.snapshot();
        assert_eq!(snapshot.get("foo", &[]), Some(vec![2]));

        let mut fork = db_clone.fork();
        fork.put("foo", vec![], vec![3]);
        db_clone.merge(fork.into_patch()).unwrap();

        let snapshot = db.snapshot();
        assert_eq!(snapshot.get("foo", &[]), Some(vec![3]));

        // Check rollback on clones
        let mut db_clone = Clone::clone(&db);
        db_clone.rollback(1);
        let snapshot = db.snapshot();
        assert_eq!(snapshot.get("foo", &[]), Some(vec![2]));
        db.rollback(1);
        let snapshot = db_clone.snapshot();
        assert_eq!(snapshot.get("foo", &[]), None);
    }

    #[test]
    fn test_checkpointdb_handler() {
        let mut db = CheckpointDb::new(MemoryDB::new());
        let mut db_handler = db.handler();

        let mut fork = db.fork();
        fork.put("foo", vec![], vec![2]);
        db.merge(fork.into_patch()).unwrap();
        let snapshot = db.snapshot();
        assert_eq!(snapshot.get("foo", &[]), Some(vec![2]));

        db_handler.rollback(1);
        let snapshot = db.snapshot();
        assert_eq!(snapshot.get("foo", &[]), None);
    }
}
<commit_msg>Update `Patch` usage<commit_after>use std::sync::{Arc, RwLock};

use exonum::storage::{Database, Patch, Result as StorageResult, Snapshot};

/// Implementation of a `Database`, which allows to rollback commits introduced by the `merge()`
/// function.
///
/// **Note:** Intended for testing purposes only. Probably inefficient.
#[derive(Clone, Debug)]
pub struct CheckpointDb<T> {
    inner: T,
    journal: Arc<RwLock<Vec<Patch>>>,
}

impl<T: Database + Clone> CheckpointDb<T> {
    /// Creates a new checkpointed database that uses the specified `db` as the underlying
    /// data storage.
    pub fn new(db: T) -> Self {
        CheckpointDb {
            inner: db,
            journal: Arc::new(RwLock::new(Vec::new())),
        }
    }

    /// Rolls this database by udoing the latest `count` `merge()` operations.
    pub fn rollback(&mut self, count: usize) -> bool {
        let journal_len = self.journal
            .read()
            .expect("Cannot acquire read lock on journal")
            .len();
        assert!(
            journal_len >= count,
            "Cannot rollback {} changes; only {} checkpoints in the journal",
            count,
            journal_len
        );

        let mut journal = self.journal.write().expect(
            "Cannot acquire write lock on journal",
        );

        for _ in 0..count {
            if let Some(patch) = journal.pop() {
                self.inner.merge(patch).expect(
                    "Cannot merge roll-back patch",
                );
            } else {
                panic!("Cannot rollback the database, the journal is empty");
            }
        }

        journal_len < count
    }

    /// Returns a handler to the database. The handler could be used to roll the database back
    /// without having the ownership to it.
    pub fn handler(&self) -> CheckpointDbHandler<T> {
        CheckpointDbHandler(Clone::clone(self))
    }
}

impl<T: Database + Clone> Database for CheckpointDb<T> {
    fn clone(&self) -> Box<Database> {
        Box::new(Clone::clone(self))
    }

    fn snapshot(&self) -> Box<Snapshot> {
        self.inner.snapshot()
    }

    fn merge(&mut self, patch: Patch) -> StorageResult<()> {
        let snapshot = self.inner.snapshot();
        self.inner.merge(patch.clone())?;
        let mut rev_fork = self.fork();

        for (name, changes) in patch {
            for (key, _) in changes {
                match snapshot.get(&name, &key) {
                    Some(value) => {
                        rev_fork.put(&name, key, value);
                    }
                    None => {
                        rev_fork.remove(&name, key);
                    }
                }
            }
        }

        {
            let mut journal = self.journal.write().expect(
                "Cannot acquire write lock on journal",
            );
            journal.push(rev_fork.into_patch());
        }
        Ok(())
    }

    fn merge_sync(&mut self, patch: Patch) -> StorageResult<()> {
        self.merge(patch)
    }
}

/// Handler to a checkpointed database.
#[derive(Clone, Debug)]
pub struct CheckpointDbHandler<T>(CheckpointDb<T>);

impl<T: Database + Clone> CheckpointDbHandler<T> {
    /// Rolls this database by udoing the latest `count` `merge()` operations.
    pub fn rollback(&mut self, count: usize) -> bool {
        self.0.rollback(count)
    }
}

#[cfg(test)]
mod tests {
    use exonum::storage::{Change, MemoryDB};
    use super::*;

    // Same as `Change`, but with trait implementation required for `Patch` comparison.
    #[derive(Debug, PartialOrd, Ord, PartialEq, Eq)]
    enum OrdChange {
        Put(Vec<u8>),
        Delete,
    }

    impl From<Change> for OrdChange {
        fn from(change: Change) -> Self {
            match change {
                Change::Put(value) => OrdChange::Put(value),
                Change::Delete => OrdChange::Delete,
            }
        }
    }

    impl<'a> From<&'a Change> for OrdChange {
        fn from(change: &'a Change) -> Self {
            match *change {
                Change::Put(ref value) => OrdChange::Put(value.clone()),
                Change::Delete => OrdChange::Delete,
            }
        }
    }

    /// Asserts that a patch contains only the specified changes.
    fn check_patch<'a, I>(patch: &Patch, changes: I)
    where
        I: IntoIterator<Item = (&'a str, Vec<u8>, Change)>,
    {
        use std::collections::BTreeSet;
        use std::iter::FromIterator;

        let mut patch_set: BTreeSet<(&str, _, _)> = BTreeSet::new();
        for (name, changes) in patch.iter() {
            for (key, value) in changes.iter() {
                patch_set.insert((name, key.clone(), OrdChange::from(value)));
            }
        }
        let expected_set = BTreeSet::from_iter(changes.into_iter().map(|(name, key, value)| {
            (name, key, OrdChange::from(value))
        }));
        assert_eq!(patch_set, expected_set);
    }

    #[test]
    fn test_checkpointdb_basics() {
        let mut db = CheckpointDb::new(MemoryDB::new());
        let mut fork = db.fork();
        fork.put("foo", vec![], vec![2]);
        db.merge(fork.into_patch()).unwrap();
        {
            let journal = db.journal.read().expect(
                "Cannot acquire read lock on journal",
            );
            assert_eq!(journal.len(), 1);
            check_patch(&journal[0], vec![("foo", vec![], Change::Delete)]);
        }

        let snapshot = db.snapshot();
        assert_eq!(snapshot.get("foo", &[]), Some(vec![2]));

        let mut fork = db.fork();
        fork.put("foo", vec![], vec![3]);
        fork.put("bar", vec![1], vec![4]);
        db.merge(fork.into_patch()).unwrap();
        {
            let journal = db.journal.read().expect(
                "Cannot acquire read lock on journal",
            );
            assert_eq!(journal.len(), 2);
            check_patch(&journal[0], vec![("foo", vec![], Change::Delete)]);
            check_patch(
                &journal[1],
                vec![("foo", vec![], Change::Put(vec![2])), ("bar", vec![1], Change::Delete)],
            );
        }

        // Check that the old snapshot still corresponds to the same DB state
        assert_eq!(snapshot.get("foo", &[]), Some(vec![2]));
        let snapshot = db.snapshot();
        assert_eq!(snapshot.get("foo", &[]), Some(vec![3]));
    }

    #[test]
    fn test_checkpointdb_rollback() {
        let mut db = CheckpointDb::new(MemoryDB::new());
        let mut fork = db.fork();
        fork.put("foo", vec![], vec![2]);
        db.merge(fork.into_patch()).unwrap();

        let mut fork = db.fork();
        fork.put("foo", vec![], vec![3]);
        fork.put("bar", vec![1], vec![4]);
        db.merge(fork.into_patch()).unwrap();
        {
            let journal = db.journal.read().expect(
                "Cannot acquire read lock on journal",
            );
            assert_eq!(journal.len(), 2);
        }

        let snapshot = db.snapshot();
        assert_eq!(snapshot.get("foo", &[]), Some(vec![3]));
        db.rollback(1);
        let snapshot = db.snapshot();
        assert_eq!(snapshot.get("foo", &[]), Some(vec![2]));
        assert_eq!(snapshot.get("bar", &[1]), None);
        {
            let journal = db.journal.read().expect(
                "Cannot acquire read lock on journal",
            );
            assert_eq!(journal.len(), 1);
        }

        // Check that DB continues working as usual after a rollback
        let mut fork = db.fork();
        fork.put("foo", vec![], vec![4]);
        fork.put("foo", vec![0, 0], vec![255]);
        db.merge(fork.into_patch()).unwrap();
        {
            let journal = db.journal.read().expect(
                "Cannot acquire read lock on journal",
            );
            assert_eq!(journal.len(), 2);
        }
        let snapshot = db.snapshot();
        assert_eq!(snapshot.get("foo", &[]), Some(vec![4]));
        assert_eq!(snapshot.get("foo", &[0, 0]), Some(vec![255]));

        let mut fork = db.fork();
        fork.put("bar", vec![1], vec![254]);
        db.merge(fork.into_patch()).unwrap();
        {
            let journal = db.journal.read().expect(
                "Cannot acquire read lock on journal",
            );
            assert_eq!(journal.len(), 3);
        }
        let new_snapshot = db.snapshot();
        assert_eq!(new_snapshot.get("foo", &[]), Some(vec![4]));
        assert_eq!(new_snapshot.get("foo", &[0, 0]), Some(vec![255]));
        assert_eq!(new_snapshot.get("bar", &[1]), Some(vec![254]));

        db.rollback(2);
        {
            let journal = db.journal.read().expect(
                "Cannot acquire read lock on journal",
            );
            assert_eq!(journal.len(), 1);
        }
        let snapshot = db.snapshot();
        assert_eq!(snapshot.get("foo", &[]), Some(vec![2]));
        assert_eq!(snapshot.get("foo", &[0, 0]), None);
        assert_eq!(snapshot.get("bar", &[1]), None);

        assert_eq!(new_snapshot.get("foo", &[]), Some(vec![4]));
        assert_eq!(new_snapshot.get("foo", &[0, 0]), Some(vec![255]));
        assert_eq!(new_snapshot.get("bar", &[1]), Some(vec![254]));
    }

    #[test]
    fn test_checkpointdb_clone() {
        let mut db = CheckpointDb::new(MemoryDB::new());
        let mut db_clone = Database::clone(&db);

        let mut fork = db.fork();
        fork.put("foo", vec![], vec![2]);
        db.merge(fork.into_patch()).unwrap();

        let snapshot = db_clone.snapshot();
        assert_eq!(snapshot.get("foo", &[]), Some(vec![2]));

        let mut fork = db_clone.fork();
        fork.put("foo", vec![], vec![3]);
        db_clone.merge(fork.into_patch()).unwrap();

        let snapshot = db.snapshot();
        assert_eq!(snapshot.get("foo", &[]), Some(vec![3]));

        // Check rollback on clones
        let mut db_clone = Clone::clone(&db);
        db_clone.rollback(1);
        let snapshot = db.snapshot();
        assert_eq!(snapshot.get("foo", &[]), Some(vec![2]));
        db.rollback(1);
        let snapshot = db_clone.snapshot();
        assert_eq!(snapshot.get("foo", &[]), None);
    }

    #[test]
    fn test_checkpointdb_handler() {
        let mut db = CheckpointDb::new(MemoryDB::new());
        let mut db_handler = db.handler();

        let mut fork = db.fork();
        fork.put("foo", vec![], vec![2]);
        db.merge(fork.into_patch()).unwrap();
        let snapshot = db.snapshot();
        assert_eq!(snapshot.get("foo", &[]), Some(vec![2]));

        db_handler.rollback(1);
        let snapshot = db.snapshot();
        assert_eq!(snapshot.get("foo", &[]), None);
    }
}
<|endoftext|>


<commit_before>use analysis::namespaces;
use case::CaseExt;
use codegen::general::{self, version_condition, version_condition_string};
use config::gobjects::GObject;
use env::Env;
use file_saver;
use library::*;
use std::collections::HashSet;
use std::io::prelude::*;
use std::io::Result;
use std::path::Path;
use traits::*;
use version::Version;

pub fn generate(env: &Env, root_path: &Path, mod_rs: &mut Vec<String>) {
    let path = root_path.join("enums.rs");
    file_saver::save_to_file(path, env.config.make_backup, |w| {
        try!(general::start_comments(w, &env.config));
        try!(writeln!(w, ""));
        try!(writeln!(w, "use ffi;"));
        if env.namespaces.glib_ns_id == namespaces::MAIN {
            try!(writeln!(w, "use error::ErrorDomain;
use translate::*;"));
        } else {
            try!(writeln!(w, "use glib_ffi;
use glib::error::ErrorDomain;
use glib::translate::*;"));
        }
        try!(writeln!(w, ""));

        let configs = env.config.objects.values()
            .filter(|c| {
                c.status.need_generate() &&
                    c.type_id.map_or(false, |tid| tid.ns_id == namespaces::MAIN)
            });
        let mut first = true;
        for config in configs {
            if let Type::Enumeration(ref enum_) = *env.library.type_(config.type_id.unwrap()) {
                if first {
                    mod_rs.push("\nmod enums;".into());
                    first = false;
                }
                if let Some (cfg) = version_condition_string(env, enum_.version, false, 0) {
                    mod_rs.push(cfg);
                }
                mod_rs.push(format!("pub use self::enums::{};", enum_.name));
                try!(generate_enum(env, w, enum_, config));
            }
        }

        Ok(())
    });
}

fn generate_enum(env: &Env, w: &mut Write, enum_: &Enumeration, config: &GObject) -> Result<()> {
    struct Member {
        name: String,
        c_name: String,
        version: Option<Version>,
    }

    let mut members: Vec<Member> = Vec::new();
    let mut vals: HashSet<String> = HashSet::new();

    for member in &enum_.members {
        let member_config = config.members.matched(&member.name);
        let is_alias = member_config.iter().any(|m| m.alias);
        if is_alias || vals.contains(&member.value) {
            continue
        }
        vals.insert(member.value.clone());
        let version = member_config.iter().filter_map(|m| m.version).next();
        members.push(Member {
            name: member.name.to_camel(),
            c_name: member.c_identifier.clone(),
            version: version,
        });
    }

    try!(version_condition(w, env, enum_.version, false, 0));
    try!(writeln!(w, "#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]"));
    try!(writeln!(w, "pub enum {} {{", enum_.name));
    for member in &members {
        try!(writeln!(w, "\t{},", member.name));
    }
    try!(writeln!(w, "{}", "    #[doc(hidden)]
    __Nonexhaustive(()),
}
"));

    try!(version_condition(w, env, enum_.version, false, 0));
    try!(writeln!(w, "#[doc(hidden)]
impl ToGlib for {name} {{
    type GlibType = ffi::{ffi_name};

    fn to_glib(&self) -> ffi::{ffi_name} {{
        match *self {{", name = enum_.name, ffi_name = enum_.c_type));
    for member in &members {
        try!(version_condition(w, env, member.version, false, 3));
        try!(writeln!(w, "\t\t\t{}::{} => ffi::{},", enum_.name, member.name, member.c_name));
    }
    try!(writeln!(w, "\t\t\t{}::__Nonexhaustive(_) => panic!(),", enum_.name));
    try!(writeln!(w, "{}",
"        }
    }
}
"));

    let assert = if env.config.generate_safety_asserts {
        "skip_assert_initialized!();\n\t\t"
    } else {
        ""
    };

    try!(version_condition(w, env, enum_.version, false, 0));
    try!(writeln!(w, "#[doc(hidden)]
impl FromGlib<ffi::{ffi_name}> for {name} {{
    fn from_glib(value: ffi::{ffi_name}) -> Self {{
        {assert}match value {{", name = enum_.name, ffi_name = enum_.c_type, assert = assert));
    for member in &members {
        try!(version_condition(w, env, member.version, false, 3));
        try!(writeln!(w, "\t\t\tffi::{} => {}::{},", member.c_name, enum_.name, member.name));
    }
    if members.len() == 1 {
        try!(writeln!(w, "\t\t\t_ => {}::__Nonexhaustive(()),", enum_.name));
    }
    try!(writeln!(w, "{}",
"        }
    }
}
"));

    if let Some(ref get_quark) = enum_.error_domain {
        let get_quark = get_quark.replace("-", "_");
        let has_failed_member = members.iter().any(|m| m.name == "Failed");

        try!(version_condition(w, env, enum_.version, false, 0));
        try!(writeln!(w, "impl ErrorDomain for {name} {{
    fn domain() -> glib_ffi::GQuark {{
        {assert}unsafe {{ ffi::{get_quark}() }}
    }}

    fn code(self) -> i32 {{
        self.to_glib() as i32
    }}

    fn from(code: i32) -> Option<Self> {{
        {assert}match code {{", name = enum_.name, get_quark = get_quark, assert = assert));

        for member in &members {
            try!(version_condition(w, env, member.version, false, 3));
            try!(writeln!(w, "\t\t\tx if x == ffi::{} as i32 => Some({}::{}),", member.c_name,
                enum_.name, member.name));
        }
        if has_failed_member {
            try!(writeln!(w, "\t\t\t_ => Some({}::Failed),", enum_.name));
        } else {
            try!(writeln!(w, "\t\t\t_ => Some({}::__Nonexhaustive(())),", enum_.name));
        }

        try!(writeln!(w, "{}",
"        }
    }
}
"));
    }

    Ok(())
}
<commit_msg>Proper ffi namespace use in glib<commit_after>use analysis::namespaces;
use case::CaseExt;
use codegen::general::{self, version_condition, version_condition_string};
use config::gobjects::GObject;
use env::Env;
use file_saver;
use library::*;
use std::collections::HashSet;
use std::io::prelude::*;
use std::io::Result;
use std::path::Path;
use traits::*;
use version::Version;

pub fn generate(env: &Env, root_path: &Path, mod_rs: &mut Vec<String>) {
    let path = root_path.join("enums.rs");
    file_saver::save_to_file(path, env.config.make_backup, |w| {
        try!(general::start_comments(w, &env.config));
        try!(writeln!(w, ""));
        try!(writeln!(w, "use ffi;"));
        if env.namespaces.glib_ns_id == namespaces::MAIN {
            try!(writeln!(w, "use ffi as glib_ffi;
use error::ErrorDomain;
use translate::*;"));
        } else {
            try!(writeln!(w, "use glib_ffi;
use glib::error::ErrorDomain;
use glib::translate::*;"));
        }
        try!(writeln!(w, ""));

        let configs = env.config.objects.values()
            .filter(|c| {
                c.status.need_generate() &&
                    c.type_id.map_or(false, |tid| tid.ns_id == namespaces::MAIN)
            });
        let mut first = true;
        for config in configs {
            if let Type::Enumeration(ref enum_) = *env.library.type_(config.type_id.unwrap()) {
                if first {
                    mod_rs.push("\nmod enums;".into());
                    first = false;
                }
                if let Some (cfg) = version_condition_string(env, enum_.version, false, 0) {
                    mod_rs.push(cfg);
                }
                mod_rs.push(format!("pub use self::enums::{};", enum_.name));
                try!(generate_enum(env, w, enum_, config));
            }
        }

        Ok(())
    });
}

fn generate_enum(env: &Env, w: &mut Write, enum_: &Enumeration, config: &GObject) -> Result<()> {
    struct Member {
        name: String,
        c_name: String,
        version: Option<Version>,
    }

    let mut members: Vec<Member> = Vec::new();
    let mut vals: HashSet<String> = HashSet::new();

    for member in &enum_.members {
        let member_config = config.members.matched(&member.name);
        let is_alias = member_config.iter().any(|m| m.alias);
        if is_alias || vals.contains(&member.value) {
            continue
        }
        vals.insert(member.value.clone());
        let version = member_config.iter().filter_map(|m| m.version).next();
        members.push(Member {
            name: member.name.to_camel(),
            c_name: member.c_identifier.clone(),
            version: version,
        });
    }

    try!(version_condition(w, env, enum_.version, false, 0));
    try!(writeln!(w, "#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]"));
    try!(writeln!(w, "pub enum {} {{", enum_.name));
    for member in &members {
        try!(writeln!(w, "\t{},", member.name));
    }
    try!(writeln!(w, "{}", "    #[doc(hidden)]
    __Nonexhaustive(()),
}
"));

    try!(version_condition(w, env, enum_.version, false, 0));
    try!(writeln!(w, "#[doc(hidden)]
impl ToGlib for {name} {{
    type GlibType = ffi::{ffi_name};

    fn to_glib(&self) -> ffi::{ffi_name} {{
        match *self {{", name = enum_.name, ffi_name = enum_.c_type));
    for member in &members {
        try!(version_condition(w, env, member.version, false, 3));
        try!(writeln!(w, "\t\t\t{}::{} => ffi::{},", enum_.name, member.name, member.c_name));
    }
    try!(writeln!(w, "\t\t\t{}::__Nonexhaustive(_) => panic!(),", enum_.name));
    try!(writeln!(w, "{}",
"        }
    }
}
"));

    let assert = if env.config.generate_safety_asserts {
        "skip_assert_initialized!();\n\t\t"
    } else {
        ""
    };

    try!(version_condition(w, env, enum_.version, false, 0));
    try!(writeln!(w, "#[doc(hidden)]
impl FromGlib<ffi::{ffi_name}> for {name} {{
    fn from_glib(value: ffi::{ffi_name}) -> Self {{
        {assert}match value {{", name = enum_.name, ffi_name = enum_.c_type, assert = assert));
    for member in &members {
        try!(version_condition(w, env, member.version, false, 3));
        try!(writeln!(w, "\t\t\tffi::{} => {}::{},", member.c_name, enum_.name, member.name));
    }
    if members.len() == 1 {
        try!(writeln!(w, "\t\t\t_ => {}::__Nonexhaustive(()),", enum_.name));
    }
    try!(writeln!(w, "{}",
"        }
    }
}
"));

    if let Some(ref get_quark) = enum_.error_domain {
        let get_quark = get_quark.replace("-", "_");
        let has_failed_member = members.iter().any(|m| m.name == "Failed");

        try!(version_condition(w, env, enum_.version, false, 0));
        try!(writeln!(w, "impl ErrorDomain for {name} {{
    fn domain() -> glib_ffi::GQuark {{
        {assert}unsafe {{ ffi::{get_quark}() }}
    }}

    fn code(self) -> i32 {{
        self.to_glib() as i32
    }}

    fn from(code: i32) -> Option<Self> {{
        {assert}match code {{", name = enum_.name, get_quark = get_quark, assert = assert));

        for member in &members {
            try!(version_condition(w, env, member.version, false, 3));
            try!(writeln!(w, "\t\t\tx if x == ffi::{} as i32 => Some({}::{}),", member.c_name,
                enum_.name, member.name));
        }
        if has_failed_member {
            try!(writeln!(w, "\t\t\t_ => Some({}::Failed),", enum_.name));
        } else {
            try!(writeln!(w, "\t\t\t_ => Some({}::__Nonexhaustive(())),", enum_.name));
        }

        try!(writeln!(w, "{}",
"        }
    }
}
"));
    }

    Ok(())
}
<|endoftext|>


<commit_before>#[cfg(test)] use std::ptr;
#[cfg(test)] use std::ffi::CStr;
#[cfg(test)] use std::os::raw::c_char;
#[cfg(test)] use std::io::Cursor;
#[cfg(test)] use llvm::prelude::*;
#[cfg(test)] use llvm::core::*;
#[cfg(test)] use llvm::execution_engine::*;
#[cfg(test)] use compileerror::{ErrorType, err, CompileError, Pos};
#[cfg(test)] use parser::{parse_module, ParseMode};
#[cfg(test)] use codegen::{CodeGenOptions, codegen, cstr};
#[cfg(test)] use ast::{TreePrinter};

#[cfg(test)]
fn run(prog: &str, dump: bool) -> Result<u64, CompileError>
{
    let mut cursor = Cursor::new(prog);
    let md = try!(parse_module(&mut cursor, "test", ParseMode::Module));
    if dump {
        md.print(0);
    }

    let opts = CodeGenOptions{
        dump_ir: dump,
        build_dir: "build".into(),
        program_name: "test".into(),
        runtime_library: "libcobraruntime.a".into(),
        optimize: true,
    };

    let mut ctx = try!(codegen(&md, &opts));

    unsafe {
        LLVMLinkInInterpreter();

        let mut ee: LLVMExecutionEngineRef = ptr::null_mut();
        let mut error_message: *mut c_char = ptr::null_mut();
        if LLVMCreateInterpreterForModule(&mut ee, ctx.take_module_ref(), &mut error_message) != 0 {
            let msg = CStr::from_ptr(error_message).to_str().expect("Invalid C string");
            let e = format!("Unable to create interpreter: {}", msg);
            LLVMDisposeMessage(error_message);
            return err(Pos::zero(), ErrorType::CodegenError(e));
        }

        let mut func: LLVMValueRef = ptr::null_mut();
        if LLVMFindFunction(ee, cstr("main"), &mut func) == 0 {
            let val = LLVMRunFunction(ee, func, 0, ptr::null_mut());
            let result = LLVMGenericValueToInt(val, 0) as u64;
            LLVMDisposeGenericValue(val);
            LLVMDisposeExecutionEngine(ee);
            Ok(result)
        } else {
            LLVMDisposeExecutionEngine(ee);
            err(Pos::zero(), ErrorType::CodegenError("No main function found".into()))
        }
    }
}

#[test]
fn test_number()
{
    assert!(run(r#"
func main() -> int:
    return 5
    "#, false).unwrap() == 5);
}

#[test]
fn test_call()
{
    assert!(run(r#"
func foo(a: int, b: int) -> int:
    return a + b * 2

func main() -> int:
    return foo(1, 2)
    "#, false).unwrap() == 5);
}

#[test]
fn test_vars()
{
    assert!(run(r#"
func main() -> int:
    const x = 7
    var y = 6
    y += x
    return y
    "#, false).unwrap() == 13);
}

#[test]
fn test_global_vars()
{
    assert!(run(r#"
const x = 7
var y = 6
func main() -> int:
    y += x
    return y
    "#, false).unwrap() == 13);
}

#[test]
fn test_struct()
{
    assert!(run(r#"
struct Point:
    var x = 0, y = 0

    pub func sum(self) -> int:
        return self.x + self.y

func main() -> int:
    var p = Point{4, 9}
    return p.sum()
    "#, false).unwrap() == 13);
}

#[test]
fn test_post_fix_increment()
{
    assert!(run(r#"
struct Foo:
    pub var z = 6

var bar = Foo{}

func main() -> int:
    var y = 6
    var foo = Foo{}
    foo.z++
    bar.z++
    y++
    return y + foo.z + bar.z
    "#, false).unwrap() == 21);
}

#[test]
fn test_arrays()
{
    assert!(run(r#"
struct Point:
    var x = 0, y = 0

    pub func sum(self) -> int:
        return self.x + self.y

func main() -> int:
    var p = [Point{4, 9}, Point{3, 4}]
    return p[0].sum() + p[1].sum()
    "#, false).unwrap() == 20);
}

#[test]
fn test_global_arrays()
{
    let v = run(r#"
struct Point:
    var x = 0, y = 0

    pub func sum(self) -> int:
        return self.x + self.y

var p = [Point{4, 9}, Point{3, 4}]

func main() -> int:
    return p[0].sum() + p[1].sum()
    "#, false).unwrap();
    assert!(v == 20);
}

#[test]
fn test_struct_with_arrays()
{
    assert!(run(r#"
struct Point:
    var coords = [0, 0]

    pub func sum(self) -> int:
        return self.coords[0] + self.coords[1]

func main() -> int:
    var p = Point{[4, 10]}
    return p.sum()
    "#, false).unwrap() == 14);
}

#[test]
fn test_array_initializer()
{
    assert!(run(r#"
func sum(v: [int]) -> int:
    var i = 0, s = 0
    while i < v.len:
        s += v[i]
        i++
    return s

func main() -> int:
    const x = 20
    return sum([x; 4])
    "#, false).unwrap() == 80);
}

#[test]
fn test_if()
{
    assert!(run(r#"
struct Point:
    var x = 0, y = 0

    pub func sum(self) -> int:
        return self.x + self.y

func main() -> int:
    const x = 77
    if x > 10:
        var p = Point{7, 9}
        return p.sum()
    else:
        return 5
    "#, false).unwrap() == 16);
}

#[test]
fn test_else()
{
    assert!(run(r#"
struct Point:
    var x = 0, y = 0

    pub func sum(self) -> int:
        return self.x + self.y

func main() -> int:
    const x = 77
    if x < 10:
        var p = Point{7, 9}
        return p.sum()
    else:
        return 5
    "#, false).unwrap() == 5);
}

#[test]
fn test_while()
{
    assert!(run(r#"
func main() -> int:
    var i = 0
    var count = 0
    while i < 10:
        count += 2
        i++
    return count
    "#, false).unwrap() == 20);
}

#[test]
fn test_slice()
{
    assert!(run(r#"
func sum(v: [int]) -> int:
    var i = 0, s = 0
    while i < v.len:
        s += v[i]
        i++
    return s

func main() -> int:
    return sum([4, 5, 6, 7])
    "#, false).unwrap() == 22);
}

/*
#[test]
fn test_traits()
{
    assert!(run(r#"
interface Sum:
    func sum(self) -> int

struct Foo:
    var a = 0, b = 0, c = 0

    func sum(self) -> int:
        return self.a + self.b + self.c

func sum(x: ~Sum) -> int:
    return x.sum()

func main() -> int:
    const f = Foo{1, 2, 3}
    return sum(f)
    "#, false).unwrap() == 6);
}
*/

#[test]
fn test_object_passing()
{
    assert!(run(r#"
struct Point:
    pub var x = 0, y = 0

func test(var p: Point) -> int:
    p.x += 1
    p.y += 1
    return p.x + p.y

func test2(var p: *Point) -> int:
    p.x += 1
    p.y += 1
    return p.x + p.y

func main() -> int:
    var a = Point{10, 12}
    var b = Point{4, 5}
    return test(a) + test2(&b) + (a.x + a.y) + (b.x + b.y)
    "#, true).unwrap() == (24 + 11 + 22 + 11));
}

/*
#[test]
fn test_object_passing_ptr()
{
    assert!(run(r#"
struct Point:
    pub var x = 0, y = 0

func test(var p: *Point) -> int:
    p.x += 1
    p.y += 1
    return p.x + p.y

func main() -> int:
    const p = Point{10, 12}
    return test(p)
    "#, false).unwrap() == 24);
}
*/
<commit_msg>Disable dumping in test_object_passing<commit_after>#[cfg(test)] use std::ptr;
#[cfg(test)] use std::ffi::CStr;
#[cfg(test)] use std::os::raw::c_char;
#[cfg(test)] use std::io::Cursor;
#[cfg(test)] use llvm::prelude::*;
#[cfg(test)] use llvm::core::*;
#[cfg(test)] use llvm::execution_engine::*;
#[cfg(test)] use compileerror::{ErrorType, err, CompileError, Pos};
#[cfg(test)] use parser::{parse_module, ParseMode};
#[cfg(test)] use codegen::{CodeGenOptions, codegen, cstr};
#[cfg(test)] use ast::{TreePrinter};

#[cfg(test)]
fn run(prog: &str, dump: bool) -> Result<u64, CompileError>
{
    let mut cursor = Cursor::new(prog);
    let md = try!(parse_module(&mut cursor, "test", ParseMode::Module));
    if dump {
        md.print(0);
    }

    let opts = CodeGenOptions{
        dump_ir: dump,
        build_dir: "build".into(),
        program_name: "test".into(),
        runtime_library: "libcobraruntime.a".into(),
        optimize: true,
    };

    let mut ctx = try!(codegen(&md, &opts));

    unsafe {
        LLVMLinkInInterpreter();

        let mut ee: LLVMExecutionEngineRef = ptr::null_mut();
        let mut error_message: *mut c_char = ptr::null_mut();
        if LLVMCreateInterpreterForModule(&mut ee, ctx.take_module_ref(), &mut error_message) != 0 {
            let msg = CStr::from_ptr(error_message).to_str().expect("Invalid C string");
            let e = format!("Unable to create interpreter: {}", msg);
            LLVMDisposeMessage(error_message);
            return err(Pos::zero(), ErrorType::CodegenError(e));
        }

        let mut func: LLVMValueRef = ptr::null_mut();
        if LLVMFindFunction(ee, cstr("main"), &mut func) == 0 {
            let val = LLVMRunFunction(ee, func, 0, ptr::null_mut());
            let result = LLVMGenericValueToInt(val, 0) as u64;
            LLVMDisposeGenericValue(val);
            LLVMDisposeExecutionEngine(ee);
            Ok(result)
        } else {
            LLVMDisposeExecutionEngine(ee);
            err(Pos::zero(), ErrorType::CodegenError("No main function found".into()))
        }
    }
}

#[test]
fn test_number()
{
    assert!(run(r#"
func main() -> int:
    return 5
    "#, false).unwrap() == 5);
}

#[test]
fn test_call()
{
    assert!(run(r#"
func foo(a: int, b: int) -> int:
    return a + b * 2

func main() -> int:
    return foo(1, 2)
    "#, false).unwrap() == 5);
}

#[test]
fn test_vars()
{
    assert!(run(r#"
func main() -> int:
    const x = 7
    var y = 6
    y += x
    return y
    "#, false).unwrap() == 13);
}

#[test]
fn test_global_vars()
{
    assert!(run(r#"
const x = 7
var y = 6
func main() -> int:
    y += x
    return y
    "#, false).unwrap() == 13);
}

#[test]
fn test_struct()
{
    assert!(run(r#"
struct Point:
    var x = 0, y = 0

    pub func sum(self) -> int:
        return self.x + self.y

func main() -> int:
    var p = Point{4, 9}
    return p.sum()
    "#, false).unwrap() == 13);
}

#[test]
fn test_post_fix_increment()
{
    assert!(run(r#"
struct Foo:
    pub var z = 6

var bar = Foo{}

func main() -> int:
    var y = 6
    var foo = Foo{}
    foo.z++
    bar.z++
    y++
    return y + foo.z + bar.z
    "#, false).unwrap() == 21);
}

#[test]
fn test_arrays()
{
    assert!(run(r#"
struct Point:
    var x = 0, y = 0

    pub func sum(self) -> int:
        return self.x + self.y

func main() -> int:
    var p = [Point{4, 9}, Point{3, 4}]
    return p[0].sum() + p[1].sum()
    "#, false).unwrap() == 20);
}

#[test]
fn test_global_arrays()
{
    let v = run(r#"
struct Point:
    var x = 0, y = 0

    pub func sum(self) -> int:
        return self.x + self.y

var p = [Point{4, 9}, Point{3, 4}]

func main() -> int:
    return p[0].sum() + p[1].sum()
    "#, false).unwrap();
    assert!(v == 20);
}

#[test]
fn test_struct_with_arrays()
{
    assert!(run(r#"
struct Point:
    var coords = [0, 0]

    pub func sum(self) -> int:
        return self.coords[0] + self.coords[1]

func main() -> int:
    var p = Point{[4, 10]}
    return p.sum()
    "#, false).unwrap() == 14);
}

#[test]
fn test_array_initializer()
{
    assert!(run(r#"
func sum(v: [int]) -> int:
    var i = 0, s = 0
    while i < v.len:
        s += v[i]
        i++
    return s

func main() -> int:
    const x = 20
    return sum([x; 4])
    "#, false).unwrap() == 80);
}

#[test]
fn test_if()
{
    assert!(run(r#"
struct Point:
    var x = 0, y = 0

    pub func sum(self) -> int:
        return self.x + self.y

func main() -> int:
    const x = 77
    if x > 10:
        var p = Point{7, 9}
        return p.sum()
    else:
        return 5
    "#, false).unwrap() == 16);
}

#[test]
fn test_else()
{
    assert!(run(r#"
struct Point:
    var x = 0, y = 0

    pub func sum(self) -> int:
        return self.x + self.y

func main() -> int:
    const x = 77
    if x < 10:
        var p = Point{7, 9}
        return p.sum()
    else:
        return 5
    "#, false).unwrap() == 5);
}

#[test]
fn test_while()
{
    assert!(run(r#"
func main() -> int:
    var i = 0
    var count = 0
    while i < 10:
        count += 2
        i++
    return count
    "#, false).unwrap() == 20);
}

#[test]
fn test_slice()
{
    assert!(run(r#"
func sum(v: [int]) -> int:
    var i = 0, s = 0
    while i < v.len:
        s += v[i]
        i++
    return s

func main() -> int:
    return sum([4, 5, 6, 7])
    "#, false).unwrap() == 22);
}

#[test]
fn test_object_passing()
{
    assert!(run(r#"
struct Point:
    pub var x = 0, y = 0

func test(var p: Point) -> int:
    p.x += 1
    p.y += 1
    return p.x + p.y

func test2(var p: *Point) -> int:
    p.x += 1
    p.y += 1
    return p.x + p.y

func main() -> int:
    var a = Point{10, 12}
    var b = Point{4, 5}
    return test(a) + test2(&b) + (a.x + a.y) + (b.x + b.y)
    "#, false).unwrap() == (24 + 11 + 22 + 11));
}

/*
#[test]
fn test_traits()
{
    assert!(run(r#"
interface Sum:
    func sum(self) -> int

struct Foo:
    var a = 0, b = 0, c = 0

    func sum(self) -> int:
        return self.a + self.b + self.c

func sum(x: ~Sum) -> int:
    return x.sum()

func main() -> int:
    const f = Foo{1, 2, 3}
    return sum(f)
    "#, false).unwrap() == 6);
}
*/
<|endoftext|>


<commit_before>use std::borrow::Cow;
use std::convert::Into;

use mio;
use mio::Token;
use mio_extras::timer::Timeout;
use url;

use io::ALL;
use message;
use protocol::CloseCode;
use result::{Error, Result};
use std::cmp::PartialEq;
use std::fmt;

#[derive(Debug, Clone)]
pub enum Signal {
    Message(message::Message),
    Close(CloseCode, Cow<'static, str>),
    Ping(Vec<u8>),
    Pong(Vec<u8>),
    Connect(url::Url),
    Shutdown,
    Timeout { delay: u64, token: Token },
    Cancel(Timeout),
}

#[derive(Debug, Clone)]
pub struct Command {
    token: Token,
    signal: Signal,
    connection_id: u32,
}

impl Command {
    pub fn token(&self) -> Token {
        self.token
    }

    pub fn into_signal(self) -> Signal {
        self.signal
    }

    pub fn connection_id(&self) -> u32 {
        self.connection_id
    }
}

/// A representation of the output of the WebSocket connection. Use this to send messages to the
/// other endpoint.
#[derive(Clone)]
pub struct Sender {
    token: Token,
    channel: mio::channel::SyncSender<Command>,
    connection_id: u32,
}

impl fmt::Debug for Sender {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f,
            "Sender {{ token: {:?}, channel: mio::channel::SyncSender<Command>, connection_id: {:?} }}",
            self.token, self.connection_id)
    }
}

impl PartialEq for Sender {
    fn eq(&self, other: &Sender) -> bool {
        self.token == other.token && self.connection_id == other.connection_id
    }
}

impl Sender {
    #[doc(hidden)]
    #[inline]
    pub fn new(
        token: Token,
        channel: mio::channel::SyncSender<Command>,
        connection_id: u32,
    ) -> Sender {
        Sender {
            token,
            channel,
            connection_id,
        }
    }

    /// A Token identifying this sender within the WebSocket.
    #[inline]
    pub fn token(&self) -> Token {
        self.token
    }

    /// A connection_id identifying this sender within the WebSocket.
    #[inline]
    pub fn connection_id(&self) -> u32 {
        self.connection_id
    }

    /// Send a message over the connection.
    #[inline]
    pub fn send<M>(&self, msg: M) -> Result<()>
    where
        M: Into<message::Message>,
    {
        self.channel
            .send(Command {
                token: self.token,
                signal: Signal::Message(msg.into()),
                connection_id: self.connection_id,
            })
            .map_err(Error::from)
    }

    /// Send a message to the endpoints of all connections.
    ///
    /// Be careful with this method. It does not discriminate between client and server connections.
    /// If your WebSocket is only functioning as a server, then usage is simple, this method will
    /// send a copy of the message to each connected client. However, if you have a WebSocket that
    /// is listening for connections and is also connected to another WebSocket, this method will
    /// broadcast a copy of the message to all the clients connected and to that WebSocket server.
    #[inline]
    pub fn broadcast<M>(&self, msg: M) -> Result<()>
    where
        M: Into<message::Message>,
    {
        self.channel
            .send(Command {
                token: ALL,
                signal: Signal::Message(msg.into()),
                connection_id: self.connection_id,
            })
            .map_err(Error::from)
    }

    /// Send a close code to the other endpoint.
    #[inline]
    pub fn close(&self, code: CloseCode) -> Result<()> {
        self.channel
            .send(Command {
                token: self.token,
                signal: Signal::Close(code, "".into()),
                connection_id: self.connection_id,
            })
            .map_err(Error::from)
    }

    /// Send a close code and provide a descriptive reason for closing.
    #[inline]
    pub fn close_with_reason<S>(&self, code: CloseCode, reason: S) -> Result<()>
    where
        S: Into<Cow<'static, str>>,
    {
        self.channel
            .send(Command {
                token: self.token,
                signal: Signal::Close(code, reason.into()),
                connection_id: self.connection_id,
            })
            .map_err(Error::from)
    }

    /// Send a ping to the other endpoint with the given test data.
    #[inline]
    pub fn ping(&self, data: Vec<u8>) -> Result<()> {
        self.channel
            .send(Command {
                token: self.token,
                signal: Signal::Ping(data),
                connection_id: self.connection_id,
            })
            .map_err(Error::from)
    }

    /// Send a pong to the other endpoint responding with the given test data.
    #[inline]
    pub fn pong(&self, data: Vec<u8>) -> Result<()> {
        self.channel
            .send(Command {
                token: self.token,
                signal: Signal::Pong(data),
                connection_id: self.connection_id,
            })
            .map_err(Error::from)
    }

    /// Queue a new connection on this WebSocket to the specified URL.
    #[inline]
    pub fn connect(&self, url: url::Url) -> Result<()> {
        self.channel
            .send(Command {
                token: self.token,
                signal: Signal::Connect(url),
                connection_id: self.connection_id,
            })
            .map_err(Error::from)
    }

    /// Request that all connections terminate and that the WebSocket stop running.
    #[inline]
    pub fn shutdown(&self) -> Result<()> {
        self.channel
            .send(Command {
                token: self.token,
                signal: Signal::Shutdown,
                connection_id: self.connection_id,
            })
            .map_err(Error::from)
    }

    /// Schedule a `token` to be sent to the WebSocket Handler's `on_timeout` method
    /// after `ms` milliseconds
    #[inline]
    pub fn timeout(&self, ms: u64, token: Token) -> Result<()> {
        self.channel
            .send(Command {
                token: self.token,
                signal: Signal::Timeout { delay: ms, token },
                connection_id: self.connection_id,
            })
            .map_err(Error::from)
    }

    /// Queue the cancellation of a previously scheduled timeout.
    ///
    /// This method is not guaranteed to prevent the timeout from occurring, because it is
    /// possible to call this method after a timeout has already occurred. It is still necessary to
    /// handle spurious timeouts.
    #[inline]
    pub fn cancel(&self, timeout: Timeout) -> Result<()> {
        self.channel
            .send(Command {
                token: self.token,
                signal: Signal::Cancel(timeout),
                connection_id: self.connection_id,
            })
            .map_err(Error::from)
    }
}
<commit_msg>Impl `Eq` and `Hash` for `Sender`.<commit_after>use std::borrow::Cow;
use std::convert::Into;

use mio;
use mio::Token;
use mio_extras::timer::Timeout;
use url;

use io::ALL;
use message;
use protocol::CloseCode;
use result::{Error, Result};
use std::cmp::PartialEq;
use std::hash::{Hash, Hasher};
use std::fmt;

#[derive(Debug, Clone)]
pub enum Signal {
    Message(message::Message),
    Close(CloseCode, Cow<'static, str>),
    Ping(Vec<u8>),
    Pong(Vec<u8>),
    Connect(url::Url),
    Shutdown,
    Timeout { delay: u64, token: Token },
    Cancel(Timeout),
}

#[derive(Debug, Clone)]
pub struct Command {
    token: Token,
    signal: Signal,
    connection_id: u32,
}

impl Command {
    pub fn token(&self) -> Token {
        self.token
    }

    pub fn into_signal(self) -> Signal {
        self.signal
    }

    pub fn connection_id(&self) -> u32 {
        self.connection_id
    }
}

/// A representation of the output of the WebSocket connection. Use this to send messages to the
/// other endpoint.
#[derive(Clone)]
pub struct Sender {
    token: Token,
    channel: mio::channel::SyncSender<Command>,
    connection_id: u32,
}

impl fmt::Debug for Sender {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f,
            "Sender {{ token: {:?}, channel: mio::channel::SyncSender<Command>, connection_id: {:?} }}",
            self.token, self.connection_id)
    }
}

impl PartialEq for Sender {
    fn eq(&self, other: &Sender) -> bool {
        self.token == other.token && self.connection_id == other.connection_id
    }
}

// Fixes #255.
impl Eq for Sender { }

// Fixes #255.
impl Hash for Sender {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.connection_id.hash(state);
        self.token.hash(state);
    }
}

impl Sender {
    #[doc(hidden)]
    #[inline]
    pub fn new(
        token: Token,
        channel: mio::channel::SyncSender<Command>,
        connection_id: u32,
    ) -> Sender {
        Sender {
            token,
            channel,
            connection_id,
        }
    }

    /// A Token identifying this sender within the WebSocket.
    #[inline]
    pub fn token(&self) -> Token {
        self.token
    }

    /// A connection_id identifying this sender within the WebSocket.
    #[inline]
    pub fn connection_id(&self) -> u32 {
        self.connection_id
    }

    /// Send a message over the connection.
    #[inline]
    pub fn send<M>(&self, msg: M) -> Result<()>
    where
        M: Into<message::Message>,
    {
        self.channel
            .send(Command {
                token: self.token,
                signal: Signal::Message(msg.into()),
                connection_id: self.connection_id,
            })
            .map_err(Error::from)
    }

    /// Send a message to the endpoints of all connections.
    ///
    /// Be careful with this method. It does not discriminate between client and server connections.
    /// If your WebSocket is only functioning as a server, then usage is simple, this method will
    /// send a copy of the message to each connected client. However, if you have a WebSocket that
    /// is listening for connections and is also connected to another WebSocket, this method will
    /// broadcast a copy of the message to all the clients connected and to that WebSocket server.
    #[inline]
    pub fn broadcast<M>(&self, msg: M) -> Result<()>
    where
        M: Into<message::Message>,
    {
        self.channel
            .send(Command {
                token: ALL,
                signal: Signal::Message(msg.into()),
                connection_id: self.connection_id,
            })
            .map_err(Error::from)
    }

    /// Send a close code to the other endpoint.
    #[inline]
    pub fn close(&self, code: CloseCode) -> Result<()> {
        self.channel
            .send(Command {
                token: self.token,
                signal: Signal::Close(code, "".into()),
                connection_id: self.connection_id,
            })
            .map_err(Error::from)
    }

    /// Send a close code and provide a descriptive reason for closing.
    #[inline]
    pub fn close_with_reason<S>(&self, code: CloseCode, reason: S) -> Result<()>
    where
        S: Into<Cow<'static, str>>,
    {
        self.channel
            .send(Command {
                token: self.token,
                signal: Signal::Close(code, reason.into()),
                connection_id: self.connection_id,
            })
            .map_err(Error::from)
    }

    /// Send a ping to the other endpoint with the given test data.
    #[inline]
    pub fn ping(&self, data: Vec<u8>) -> Result<()> {
        self.channel
            .send(Command {
                token: self.token,
                signal: Signal::Ping(data),
                connection_id: self.connection_id,
            })
            .map_err(Error::from)
    }

    /// Send a pong to the other endpoint responding with the given test data.
    #[inline]
    pub fn pong(&self, data: Vec<u8>) -> Result<()> {
        self.channel
            .send(Command {
                token: self.token,
                signal: Signal::Pong(data),
                connection_id: self.connection_id,
            })
            .map_err(Error::from)
    }

    /// Queue a new connection on this WebSocket to the specified URL.
    #[inline]
    pub fn connect(&self, url: url::Url) -> Result<()> {
        self.channel
            .send(Command {
                token: self.token,
                signal: Signal::Connect(url),
                connection_id: self.connection_id,
            })
            .map_err(Error::from)
    }

    /// Request that all connections terminate and that the WebSocket stop running.
    #[inline]
    pub fn shutdown(&self) -> Result<()> {
        self.channel
            .send(Command {
                token: self.token,
                signal: Signal::Shutdown,
                connection_id: self.connection_id,
            })
            .map_err(Error::from)
    }

    /// Schedule a `token` to be sent to the WebSocket Handler's `on_timeout` method
    /// after `ms` milliseconds
    #[inline]
    pub fn timeout(&self, ms: u64, token: Token) -> Result<()> {
        self.channel
            .send(Command {
                token: self.token,
                signal: Signal::Timeout { delay: ms, token },
                connection_id: self.connection_id,
            })
            .map_err(Error::from)
    }

    /// Queue the cancellation of a previously scheduled timeout.
    ///
    /// This method is not guaranteed to prevent the timeout from occurring, because it is
    /// possible to call this method after a timeout has already occurred. It is still necessary to
    /// handle spurious timeouts.
    #[inline]
    pub fn cancel(&self, timeout: Timeout) -> Result<()> {
        self.channel
            .send(Command {
                token: self.token,
                signal: Signal::Cancel(timeout),
                connection_id: self.connection_id,
            })
            .map_err(Error::from)
    }
}
<|endoftext|>


<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{account_address::AccountAddress, event::EventKey};
use anyhow::Result;
use move_core_types::move_resource::MoveResource;
use serde::{Deserialize, Serialize};

/// Struct that represents a NewBlockEvent.
#[derive(Debug, Serialize, Deserialize)]
pub struct NewBlockEvent {
    round: u64,
    proposer: AccountAddress,
    previous_block_votes: Vec<AccountAddress>,
    time_micro_seconds: u64,
}

impl NewBlockEvent {
    pub fn round(&self) -> u64 {
        self.round
    }

    pub fn proposer(&self) -> AccountAddress {
        self.proposer
    }

    pub fn previous_block_votes(&self) -> &[AccountAddress] {
        &self.previous_block_votes
    }

    pub fn proposed_time(&self) -> u64 {
        self.time_micro_seconds
    }

    pub fn try_from_bytes(bytes: &[u8]) -> Result<Self> {
        lcs::from_bytes(bytes).map_err(Into::into)
    }

    pub fn event_key() -> EventKey {
        crate::block_metadata::new_block_event_key()
    }
}

impl MoveResource for NewBlockEvent {
    const MODULE_NAME: &'static str = "LibraBlock";
    const STRUCT_NAME: &'static str = "NewBlockEvent";
}
<commit_msg>Removing previous_block_votes from types/src/account_config/events/new_block.rs<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{account_address::AccountAddress, event::EventKey};
use anyhow::Result;
use move_core_types::move_resource::MoveResource;
use serde::{Deserialize, Serialize};

/// Struct that represents a NewBlockEvent.
#[derive(Debug, Serialize, Deserialize)]
pub struct NewBlockEvent {
    round: u64,
    proposer: AccountAddress,
    previous_block_votes: Vec<AccountAddress>,
    time_micro_seconds: u64,
}

impl NewBlockEvent {
    pub fn round(&self) -> u64 {
        self.round
    }

    pub fn proposer(&self) -> AccountAddress {
        self.proposer
    }

    pub fn proposed_time(&self) -> u64 {
        self.time_micro_seconds
    }

    pub fn try_from_bytes(bytes: &[u8]) -> Result<Self> {
        lcs::from_bytes(bytes).map_err(Into::into)
    }

    pub fn event_key() -> EventKey {
        crate::block_metadata::new_block_event_key()
    }
}

impl MoveResource for NewBlockEvent {
    const MODULE_NAME: &'static str = "LibraBlock";
    const STRUCT_NAME: &'static str = "NewBlockEvent";
}
<|endoftext|>