<commit_before>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! CoalesceBatchesExec combines small batches into larger batches for more efficient use of
//! vectorized processing by upstream operators.

use std::any::Any;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

use crate::error::{DataFusionError, Result};
use crate::physical_plan::{
    ExecutionPlan, Partitioning, RecordBatchStream, SendableRecordBatchStream,
};

use arrow::compute::kernels::concat::concat;
use arrow::datatypes::SchemaRef;
use arrow::error::Result as ArrowResult;
use arrow::record_batch::RecordBatch;
use async_trait::async_trait;
use futures::stream::{Stream, StreamExt};
use log::debug;

/// CoalesceBatchesExec combines small batches into larger batches for more efficient use of
/// vectorized processing by upstream operators.
#[derive(Debug)]
pub struct CoalesceBatchesExec {
    /// The input plan
    input: Arc<dyn ExecutionPlan>,
    /// Minimum number of rows for coalesces batches
    target_batch_size: usize,
}

impl CoalesceBatchesExec {
    /// Create a new CoalesceBatchesExec
    pub fn new(input: Arc<dyn ExecutionPlan>, target_batch_size: usize) -> Self {
        Self {
            input,
            target_batch_size,
        }
    }
}

#[async_trait]
impl ExecutionPlan for CoalesceBatchesExec {
    /// Return a reference to Any that can be used for downcasting
    fn as_any(&self) -> &dyn Any {
        self
    }

    /// Get the schema for this execution plan
    fn schema(&self) -> SchemaRef {
        // The coalesce batches operator does not make any changes to the schema of its input
        self.input.schema()
    }

    fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
        vec![self.input.clone()]
    }

    /// Get the output partitioning of this plan
    fn output_partitioning(&self) -> Partitioning {
        // The coalesce batches operator does not make any changes to the partitioning of its input
        self.input.output_partitioning()
    }

    fn with_new_children(
        &self,
        children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        match children.len() {
            1 => Ok(Arc::new(CoalesceBatchesExec::new(
                children[0].clone(),
                self.target_batch_size,
            ))),
            _ => Err(DataFusionError::Internal(
                "CoalesceBatchesExec wrong number of children".to_string(),
            )),
        }
    }

    async fn execute(&self, partition: usize) -> Result<SendableRecordBatchStream> {
        Ok(Box::pin(CoalesceBatchesStream {
            input: self.input.execute(partition).await?,
            schema: self.input.schema().clone(),
            target_batch_size: self.target_batch_size.clone(),
            buffer: Vec::new(),
            buffered_rows: 0,
        }))
    }
}

struct CoalesceBatchesStream {
    /// The input plan
    input: SendableRecordBatchStream,
    /// The input schema
    schema: SchemaRef,
    /// Minimum number of rows for coalesces batches
    target_batch_size: usize,
    /// Buffered batches
    buffer: Vec<RecordBatch>,
    /// Buffered row count
    buffered_rows: usize,
}

impl Stream for CoalesceBatchesStream {
    type Item = ArrowResult<RecordBatch>;

    fn poll_next(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Self::Item>> {
        loop {
            let input_batch = self.input.poll_next_unpin(cx);
            match input_batch {
                Poll::Ready(x) => match x {
                    Some(Ok(ref batch)) => {
                        if batch.num_rows() >= self.target_batch_size
                            && self.buffer.is_empty()
                        {
                            return Poll::Ready(Some(Ok(batch.clone())));
                        } else if batch.num_rows() == 0 {
                            // discard empty batches
                        } else {
                            // add to the buffered batches
                            self.buffer.push(batch.clone());
                            self.buffered_rows += batch.num_rows();
                            // check to see if we have enough batches yet
                            if self.buffered_rows >= self.target_batch_size {
                                // combine the batches and return
                                let batch = concat_batches(
                                    &self.schema,
                                    &self.buffer,
                                    self.buffered_rows,
                                )?;
                                // reset buffer state
                                self.buffer.clear();
                                self.buffered_rows = 0;
                                // return batch
                                return Poll::Ready(Some(Ok(batch)));
                            }
                        }
                    }
                    None => {
                        // we have reached the end of the input stream but there could still
                        // be buffered batches
                        if self.buffer.is_empty() {
                            return Poll::Ready(None);
                        } else {
                            // combine the batches and return
                            let batch = concat_batches(
                                &self.schema,
                                &self.buffer,
                                self.buffered_rows,
                            )?;
                            // reset buffer state
                            self.buffer.clear();
                            self.buffered_rows = 0;
                            // return batch
                            return Poll::Ready(Some(Ok(batch)));
                        }
                    }
                    other => return Poll::Ready(other),
                },
                Poll::Pending => return Poll::Pending,
            }
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        // we can't predict the size of incoming batches so re-use the size hint from the input
        self.input.size_hint()
    }
}

impl RecordBatchStream for CoalesceBatchesStream {
    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }
}

fn concat_batches(
    schema: &SchemaRef,
    batches: &[RecordBatch],
    row_count: usize,
) -> ArrowResult<RecordBatch> {
    let mut arrays = Vec::with_capacity(schema.fields().len());
    for i in 0..schema.fields().len() {
        let array = concat(
            &batches
                .iter()
                .map(|batch| batch.column(i).as_ref())
                .collect::<Vec<_>>(),
        )?;
        arrays.push(array);
    }
    debug!(
        "Combined {} batches containing {} rows",
        batches.len(),
        row_count
    );
    RecordBatch::try_new(schema.clone(), arrays)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::physical_plan::memory::MemoryExec;
    use arrow::array::UInt32Array;
    use arrow::datatypes::{DataType, Field, Schema};

    #[tokio::test(threaded_scheduler)]
    async fn test_concat_batches() -> Result<()> {
        let schema = test_schema();
        let partition = create_vec_batches(&schema, 10)?;
        let partitions = vec![partition];

        let output_partitions = coalesce_batches(&schema, partitions, 20).await?;
        assert_eq!(1, output_partitions.len());

        // input is 10 batches x 8 rows (80 rows)
        // expected output is batches of at least 20 rows (except for the final batch)
        let batches = &output_partitions[0];
        assert_eq!(4, batches.len());
        assert_eq!(24, batches[0].num_rows());
        assert_eq!(24, batches[1].num_rows());
        assert_eq!(24, batches[2].num_rows());
        assert_eq!(8, batches[3].num_rows());

        Ok(())
    }

    fn test_schema() -> Arc<Schema> {
        Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)]))
    }

    fn create_vec_batches(
        schema: &Arc<Schema>,
        num_batches: usize,
    ) -> Result<Vec<RecordBatch>> {
        let batch = create_batch(schema);
        let mut vec = Vec::with_capacity(num_batches);
        for _ in 0..num_batches {
            vec.push(batch.clone());
        }
        Ok(vec)
    }

    fn create_batch(schema: &Arc<Schema>) -> RecordBatch {
        RecordBatch::try_new(
            schema.clone(),
            vec![Arc::new(UInt32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8]))],
        )
        .unwrap()
    }

    async fn coalesce_batches(
        schema: &SchemaRef,
        input_partitions: Vec<Vec<RecordBatch>>,
        target_batch_size: usize,
    ) -> Result<Vec<Vec<RecordBatch>>> {
        // create physical plan
        let exec = MemoryExec::try_new(&input_partitions, schema.clone(), None)?;
        let exec: Arc<dyn ExecutionPlan> =
            Arc::new(CoalesceBatchesExec::new(Arc::new(exec), target_batch_size));

        // execute and collect results
        let output_partition_count = exec.output_partitioning().partition_count();
        let mut output_partitions = Vec::with_capacity(output_partition_count);
        for i in 0..output_partition_count {
            // execute this *output* partition and collect all batches
            let mut stream = exec.execute(i).await?;
            let mut batches = vec![];
            while let Some(result) = stream.next().await {
                batches.push(result?);
            }
            output_partitions.push(batches);
        }
        Ok(output_partitions)
    }
}
<commit_msg>ARROW-11084: [Rust] Fixed clippy<commit_after>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! CoalesceBatchesExec combines small batches into larger batches for more efficient use of
//! vectorized processing by upstream operators.

use std::any::Any;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

use crate::error::{DataFusionError, Result};
use crate::physical_plan::{
    ExecutionPlan, Partitioning, RecordBatchStream, SendableRecordBatchStream,
};

use arrow::compute::kernels::concat::concat;
use arrow::datatypes::SchemaRef;
use arrow::error::Result as ArrowResult;
use arrow::record_batch::RecordBatch;
use async_trait::async_trait;
use futures::stream::{Stream, StreamExt};
use log::debug;

/// CoalesceBatchesExec combines small batches into larger batches for more efficient use of
/// vectorized processing by upstream operators.
#[derive(Debug)]
pub struct CoalesceBatchesExec {
    /// The input plan
    input: Arc<dyn ExecutionPlan>,
    /// Minimum number of rows for coalesces batches
    target_batch_size: usize,
}

impl CoalesceBatchesExec {
    /// Create a new CoalesceBatchesExec
    pub fn new(input: Arc<dyn ExecutionPlan>, target_batch_size: usize) -> Self {
        Self {
            input,
            target_batch_size,
        }
    }
}

#[async_trait]
impl ExecutionPlan for CoalesceBatchesExec {
    /// Return a reference to Any that can be used for downcasting
    fn as_any(&self) -> &dyn Any {
        self
    }

    /// Get the schema for this execution plan
    fn schema(&self) -> SchemaRef {
        // The coalesce batches operator does not make any changes to the schema of its input
        self.input.schema()
    }

    fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
        vec![self.input.clone()]
    }

    /// Get the output partitioning of this plan
    fn output_partitioning(&self) -> Partitioning {
        // The coalesce batches operator does not make any changes to the partitioning of its input
        self.input.output_partitioning()
    }

    fn with_new_children(
        &self,
        children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        match children.len() {
            1 => Ok(Arc::new(CoalesceBatchesExec::new(
                children[0].clone(),
                self.target_batch_size,
            ))),
            _ => Err(DataFusionError::Internal(
                "CoalesceBatchesExec wrong number of children".to_string(),
            )),
        }
    }

    async fn execute(&self, partition: usize) -> Result<SendableRecordBatchStream> {
        Ok(Box::pin(CoalesceBatchesStream {
            input: self.input.execute(partition).await?,
            schema: self.input.schema(),
            target_batch_size: self.target_batch_size,
            buffer: Vec::new(),
            buffered_rows: 0,
        }))
    }
}

struct CoalesceBatchesStream {
    /// The input plan
    input: SendableRecordBatchStream,
    /// The input schema
    schema: SchemaRef,
    /// Minimum number of rows for coalesces batches
    target_batch_size: usize,
    /// Buffered batches
    buffer: Vec<RecordBatch>,
    /// Buffered row count
    buffered_rows: usize,
}

impl Stream for CoalesceBatchesStream {
    type Item = ArrowResult<RecordBatch>;

    fn poll_next(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Self::Item>> {
        loop {
            let input_batch = self.input.poll_next_unpin(cx);
            match input_batch {
                Poll::Ready(x) => match x {
                    Some(Ok(ref batch)) => {
                        if batch.num_rows() >= self.target_batch_size
                            && self.buffer.is_empty()
                        {
                            return Poll::Ready(Some(Ok(batch.clone())));
                        } else if batch.num_rows() == 0 {
                            // discard empty batches
                        } else {
                            // add to the buffered batches
                            self.buffer.push(batch.clone());
                            self.buffered_rows += batch.num_rows();
                            // check to see if we have enough batches yet
                            if self.buffered_rows >= self.target_batch_size {
                                // combine the batches and return
                                let batch = concat_batches(
                                    &self.schema,
                                    &self.buffer,
                                    self.buffered_rows,
                                )?;
                                // reset buffer state
                                self.buffer.clear();
                                self.buffered_rows = 0;
                                // return batch
                                return Poll::Ready(Some(Ok(batch)));
                            }
                        }
                    }
                    None => {
                        // we have reached the end of the input stream but there could still
                        // be buffered batches
                        if self.buffer.is_empty() {
                            return Poll::Ready(None);
                        } else {
                            // combine the batches and return
                            let batch = concat_batches(
                                &self.schema,
                                &self.buffer,
                                self.buffered_rows,
                            )?;
                            // reset buffer state
                            self.buffer.clear();
                            self.buffered_rows = 0;
                            // return batch
                            return Poll::Ready(Some(Ok(batch)));
                        }
                    }
                    other => return Poll::Ready(other),
                },
                Poll::Pending => return Poll::Pending,
            }
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        // we can't predict the size of incoming batches so re-use the size hint from the input
        self.input.size_hint()
    }
}

impl RecordBatchStream for CoalesceBatchesStream {
    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }
}

fn concat_batches(
    schema: &SchemaRef,
    batches: &[RecordBatch],
    row_count: usize,
) -> ArrowResult<RecordBatch> {
    let mut arrays = Vec::with_capacity(schema.fields().len());
    for i in 0..schema.fields().len() {
        let array = concat(
            &batches
                .iter()
                .map(|batch| batch.column(i).as_ref())
                .collect::<Vec<_>>(),
        )?;
        arrays.push(array);
    }
    debug!(
        "Combined {} batches containing {} rows",
        batches.len(),
        row_count
    );
    RecordBatch::try_new(schema.clone(), arrays)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::physical_plan::memory::MemoryExec;
    use arrow::array::UInt32Array;
    use arrow::datatypes::{DataType, Field, Schema};

    #[tokio::test(threaded_scheduler)]
    async fn test_concat_batches() -> Result<()> {
        let schema = test_schema();
        let partition = create_vec_batches(&schema, 10)?;
        let partitions = vec![partition];

        let output_partitions = coalesce_batches(&schema, partitions, 20).await?;
        assert_eq!(1, output_partitions.len());

        // input is 10 batches x 8 rows (80 rows)
        // expected output is batches of at least 20 rows (except for the final batch)
        let batches = &output_partitions[0];
        assert_eq!(4, batches.len());
        assert_eq!(24, batches[0].num_rows());
        assert_eq!(24, batches[1].num_rows());
        assert_eq!(24, batches[2].num_rows());
        assert_eq!(8, batches[3].num_rows());

        Ok(())
    }

    fn test_schema() -> Arc<Schema> {
        Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)]))
    }

    fn create_vec_batches(
        schema: &Arc<Schema>,
        num_batches: usize,
    ) -> Result<Vec<RecordBatch>> {
        let batch = create_batch(schema);
        let mut vec = Vec::with_capacity(num_batches);
        for _ in 0..num_batches {
            vec.push(batch.clone());
        }
        Ok(vec)
    }

    fn create_batch(schema: &Arc<Schema>) -> RecordBatch {
        RecordBatch::try_new(
            schema.clone(),
            vec![Arc::new(UInt32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8]))],
        )
        .unwrap()
    }

    async fn coalesce_batches(
        schema: &SchemaRef,
        input_partitions: Vec<Vec<RecordBatch>>,
        target_batch_size: usize,
    ) -> Result<Vec<Vec<RecordBatch>>> {
        // create physical plan
        let exec = MemoryExec::try_new(&input_partitions, schema.clone(), None)?;
        let exec: Arc<dyn ExecutionPlan> =
            Arc::new(CoalesceBatchesExec::new(Arc::new(exec), target_batch_size));

        // execute and collect results
        let output_partition_count = exec.output_partitioning().partition_count();
        let mut output_partitions = Vec::with_capacity(output_partition_count);
        for i in 0..output_partition_count {
            // execute this *output* partition and collect all batches
            let mut stream = exec.execute(i).await?;
            let mut batches = vec![];
            while let Some(result) = stream.next().await {
                batches.push(result?);
            }
            output_partitions.push(batches);
        }
        Ok(output_partitions)
    }
}
<|endoftext|>


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

use crate::{corpus_from_strategy, fuzz_data_to_value, FuzzTargetImpl};
use accumulator::test_helpers::{
    arb_hash_batch, arb_list_of_hash_batches, arb_three_hash_batches, arb_two_hash_batches,
    test_append_empty_impl, test_append_many_impl, test_consistency_proof_impl,
    test_get_frozen_subtree_hashes_impl, test_proof_impl, test_range_proof_impl,
};
use diem_jellyfish_merkle::test_helper::{
    arb_existent_kvs_and_nonexistent_keys, arb_kv_pair_with_distinct_last_nibble,
    arb_tree_with_index, test_get_range_proof, test_get_with_proof,
    test_get_with_proof_with_distinct_last_nibble,
};
use diem_proptest_helpers::ValueGenerator;
use diem_types::account_state_blob::AccountStateBlob;
use diemdb::{
    schema::fuzzing::fuzz_decode, test_helper::arb_blocks_to_commit, test_save_blocks_impl,
};
use proptest::{collection::vec, prelude::*};

#[derive(Clone, Debug, Default)]
pub struct StorageSaveBlocks;

impl FuzzTargetImpl for StorageSaveBlocks {
    fn description(&self) -> &'static str {
        "Storage save blocks"
    }

    fn generate(&self, _idx: usize, _gen: &mut ValueGenerator) -> Option<Vec<u8>> {
        Some(corpus_from_strategy(arb_blocks_to_commit()))
    }

    fn fuzz(&self, data: &[u8]) {
        let input = fuzz_data_to_value(data, arb_blocks_to_commit());
        test_save_blocks_impl(input);
    }
}

#[derive(Clone, Debug, Default)]
pub struct StorageSchemaDecode;

impl FuzzTargetImpl for StorageSchemaDecode {
    fn description(&self) -> &'static str {
        "Storage schemas do not panic on corrupted bytes."
    }

    fn generate(&self, _idx: usize, gen: &mut ValueGenerator) -> Option<Vec<u8>> {
        Some(gen.generate(prop_oneof![
            100 => vec(any::<u8>(), 0..1024),
            1 => vec(any::<u8>(), 1024..1024 * 10),
        ]))
    }

    fn fuzz(&self, data: &[u8]) {
        fuzz_decode(data)
    }
}

//============== JellyfishMerkleTree =============

#[derive(Clone, Debug, Default)]
pub struct JellyfishGetWithProof;

impl FuzzTargetImpl for JellyfishGetWithProof {
    fn description(&self) -> &'static str {
        "JellyfishMerkleTree get with proof"
    }

    fn generate(&self, _idx: usize, _gen: &mut ValueGenerator) -> Option<Vec<u8>> {
        Some(corpus_from_strategy(
            arb_existent_kvs_and_nonexistent_keys::<AccountStateBlob>(1000, 100),
        ))
    }

    fn fuzz(&self, data: &[u8]) {
        let input = fuzz_data_to_value(
            data,
            arb_existent_kvs_and_nonexistent_keys::<AccountStateBlob>(1000, 100),
        );
        test_get_with_proof(input);
    }
}

#[derive(Clone, Debug, Default)]
pub struct JellyfishGetWithProofWithDistinctLastNibble;

impl FuzzTargetImpl for JellyfishGetWithProofWithDistinctLastNibble {
    fn description(&self) -> &'static str {
        "JellyfishMerkleTree get with proof with distinct last nibble"
    }

    fn generate(&self, _idx: usize, _gen: &mut ValueGenerator) -> Option<Vec<u8>> {
        Some(corpus_from_strategy(
            arb_kv_pair_with_distinct_last_nibble::<AccountStateBlob>(),
        ))
    }

    fn fuzz(&self, data: &[u8]) {
        let input = fuzz_data_to_value(
            data,
            arb_kv_pair_with_distinct_last_nibble::<AccountStateBlob>(),
        );
        test_get_with_proof_with_distinct_last_nibble(input);
    }
}

#[derive(Clone, Debug, Default)]
pub struct JellyfishGetRangeProof;

impl FuzzTargetImpl for JellyfishGetRangeProof {
    fn description(&self) -> &'static str {
        "JellyfishMerkleTree get range proof"
    }

    fn generate(&self, _idx: usize, _gen: &mut ValueGenerator) -> Option<Vec<u8>> {
        Some(corpus_from_strategy(
            arb_tree_with_index::<AccountStateBlob>(1000),
        ))
    }

    fn fuzz(&self, data: &[u8]) {
        let input = fuzz_data_to_value(data, arb_tree_with_index::<AccountStateBlob>(1000));
        test_get_range_proof(input);
    }
}

#[derive(Clone, Debug, Default)]
pub struct AccumulatorFrozenSubtreeHashes;

impl FuzzTargetImpl for AccumulatorFrozenSubtreeHashes {
    fn description(&self) -> &'static str {
        "Accumulator frozen subtree hashes"
    }

    fn generate(&self, _idx: usize, _gen: &mut ValueGenerator) -> Option<Vec<u8>> {
        Some(corpus_from_strategy(arb_hash_batch(1000)))
    }

    fn fuzz(&self, data: &[u8]) {
        let input = fuzz_data_to_value(data, arb_hash_batch(1000));
        test_get_frozen_subtree_hashes_impl(input);
    }
}

//============== Accumulator =============

#[derive(Clone, Debug, Default)]
pub struct AccumulatorProof;

impl FuzzTargetImpl for AccumulatorProof {
    fn description(&self) -> &'static str {
        "Accumulator proof"
    }

    fn generate(&self, _idx: usize, _gen: &mut ValueGenerator) -> Option<Vec<u8>> {
        Some(corpus_from_strategy(arb_two_hash_batches(100)))
    }

    fn fuzz(&self, data: &[u8]) {
        let input = fuzz_data_to_value(data, arb_two_hash_batches(100));
        test_proof_impl(input);
    }
}

#[derive(Clone, Debug, Default)]
pub struct AccumulatorConsistencyProof;

impl FuzzTargetImpl for AccumulatorConsistencyProof {
    fn description(&self) -> &'static str {
        "Accumulator consistency proof"
    }

    fn generate(&self, _idx: usize, _gen: &mut ValueGenerator) -> Option<Vec<u8>> {
        Some(corpus_from_strategy(arb_two_hash_batches(100)))
    }

    fn fuzz(&self, data: &[u8]) {
        let input = fuzz_data_to_value(data, arb_two_hash_batches(100));
        test_consistency_proof_impl(input);
    }
}

#[derive(Clone, Debug, Default)]
pub struct AccumulatorRangeProof;

impl FuzzTargetImpl for AccumulatorRangeProof {
    fn description(&self) -> &'static str {
        "Accumulator range proof"
    }

    fn generate(&self, _idx: usize, _gen: &mut ValueGenerator) -> Option<Vec<u8>> {
        Some(corpus_from_strategy(arb_three_hash_batches(100)))
    }

    fn fuzz(&self, data: &[u8]) {
        let input = fuzz_data_to_value(data, arb_three_hash_batches(100));
        test_range_proof_impl(input);
    }
}

#[derive(Clone, Debug, Default)]
pub struct AccumulatorAppendMany;

impl FuzzTargetImpl for AccumulatorAppendMany {
    fn description(&self) -> &'static str {
        "Accumulator amend many leaves"
    }

    fn generate(&self, _idx: usize, _gen: &mut ValueGenerator) -> Option<Vec<u8>> {
        Some(corpus_from_strategy(arb_list_of_hash_batches(10, 10)))
    }

    fn fuzz(&self, data: &[u8]) {
        let input = fuzz_data_to_value(data, arb_list_of_hash_batches(10, 10));
        test_append_many_impl(input);
    }
}

#[derive(Clone, Debug, Default)]
pub struct AccumulatorAppendEmpty;

impl FuzzTargetImpl for AccumulatorAppendEmpty {
    fn description(&self) -> &'static str {
        "Accumulator amend empty leave"
    }

    fn generate(&self, _idx: usize, _gen: &mut ValueGenerator) -> Option<Vec<u8>> {
        Some(corpus_from_strategy(arb_hash_batch(100)))
    }

    fn fuzz(&self, data: &[u8]) {
        let input = fuzz_data_to_value(data, arb_hash_batch(100));
        test_append_empty_impl(input);
    }
}
<commit_msg>[storage] make fuzzer JellyfishGetRangeProof faster<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{corpus_from_strategy, fuzz_data_to_value, FuzzTargetImpl};
use accumulator::test_helpers::{
    arb_hash_batch, arb_list_of_hash_batches, arb_three_hash_batches, arb_two_hash_batches,
    test_append_empty_impl, test_append_many_impl, test_consistency_proof_impl,
    test_get_frozen_subtree_hashes_impl, test_proof_impl, test_range_proof_impl,
};
use diem_jellyfish_merkle::test_helper::{
    arb_existent_kvs_and_nonexistent_keys, arb_kv_pair_with_distinct_last_nibble,
    arb_tree_with_index, test_get_range_proof, test_get_with_proof,
    test_get_with_proof_with_distinct_last_nibble,
};
use diem_proptest_helpers::ValueGenerator;
use diem_types::account_state_blob::AccountStateBlob;
use diemdb::{
    schema::fuzzing::fuzz_decode, test_helper::arb_blocks_to_commit, test_save_blocks_impl,
};
use proptest::{collection::vec, prelude::*};

#[derive(Clone, Debug, Default)]
pub struct StorageSaveBlocks;

impl FuzzTargetImpl for StorageSaveBlocks {
    fn description(&self) -> &'static str {
        "Storage save blocks"
    }

    fn generate(&self, _idx: usize, _gen: &mut ValueGenerator) -> Option<Vec<u8>> {
        Some(corpus_from_strategy(arb_blocks_to_commit()))
    }

    fn fuzz(&self, data: &[u8]) {
        let input = fuzz_data_to_value(data, arb_blocks_to_commit());
        test_save_blocks_impl(input);
    }
}

#[derive(Clone, Debug, Default)]
pub struct StorageSchemaDecode;

impl FuzzTargetImpl for StorageSchemaDecode {
    fn description(&self) -> &'static str {
        "Storage schemas do not panic on corrupted bytes."
    }

    fn generate(&self, _idx: usize, gen: &mut ValueGenerator) -> Option<Vec<u8>> {
        Some(gen.generate(prop_oneof![
            100 => vec(any::<u8>(), 0..1024),
            1 => vec(any::<u8>(), 1024..1024 * 10),
        ]))
    }

    fn fuzz(&self, data: &[u8]) {
        fuzz_decode(data)
    }
}

//============== JellyfishMerkleTree =============

#[derive(Clone, Debug, Default)]
pub struct JellyfishGetWithProof;

impl FuzzTargetImpl for JellyfishGetWithProof {
    fn description(&self) -> &'static str {
        "JellyfishMerkleTree get with proof"
    }

    fn generate(&self, _idx: usize, _gen: &mut ValueGenerator) -> Option<Vec<u8>> {
        Some(corpus_from_strategy(
            arb_existent_kvs_and_nonexistent_keys::<AccountStateBlob>(1000, 100),
        ))
    }

    fn fuzz(&self, data: &[u8]) {
        let input = fuzz_data_to_value(
            data,
            arb_existent_kvs_and_nonexistent_keys::<AccountStateBlob>(1000, 100),
        );
        test_get_with_proof(input);
    }
}

#[derive(Clone, Debug, Default)]
pub struct JellyfishGetWithProofWithDistinctLastNibble;

impl FuzzTargetImpl for JellyfishGetWithProofWithDistinctLastNibble {
    fn description(&self) -> &'static str {
        "JellyfishMerkleTree get with proof with distinct last nibble"
    }

    fn generate(&self, _idx: usize, _gen: &mut ValueGenerator) -> Option<Vec<u8>> {
        Some(corpus_from_strategy(
            arb_kv_pair_with_distinct_last_nibble::<AccountStateBlob>(),
        ))
    }

    fn fuzz(&self, data: &[u8]) {
        let input = fuzz_data_to_value(
            data,
            arb_kv_pair_with_distinct_last_nibble::<AccountStateBlob>(),
        );
        test_get_with_proof_with_distinct_last_nibble(input);
    }
}

#[derive(Clone, Debug, Default)]
pub struct JellyfishGetRangeProof;

impl FuzzTargetImpl for JellyfishGetRangeProof {
    fn description(&self) -> &'static str {
        "JellyfishMerkleTree get range proof"
    }

    fn generate(&self, _idx: usize, _gen: &mut ValueGenerator) -> Option<Vec<u8>> {
        Some(corpus_from_strategy(
            arb_tree_with_index::<AccountStateBlob>(100),
        ))
    }

    fn fuzz(&self, data: &[u8]) {
        let input = fuzz_data_to_value(data, arb_tree_with_index::<AccountStateBlob>(100));
        test_get_range_proof(input);
    }
}

#[derive(Clone, Debug, Default)]
pub struct AccumulatorFrozenSubtreeHashes;

impl FuzzTargetImpl for AccumulatorFrozenSubtreeHashes {
    fn description(&self) -> &'static str {
        "Accumulator frozen subtree hashes"
    }

    fn generate(&self, _idx: usize, _gen: &mut ValueGenerator) -> Option<Vec<u8>> {
        Some(corpus_from_strategy(arb_hash_batch(1000)))
    }

    fn fuzz(&self, data: &[u8]) {
        let input = fuzz_data_to_value(data, arb_hash_batch(1000));
        test_get_frozen_subtree_hashes_impl(input);
    }
}

//============== Accumulator =============

#[derive(Clone, Debug, Default)]
pub struct AccumulatorProof;

impl FuzzTargetImpl for AccumulatorProof {
    fn description(&self) -> &'static str {
        "Accumulator proof"
    }

    fn generate(&self, _idx: usize, _gen: &mut ValueGenerator) -> Option<Vec<u8>> {
        Some(corpus_from_strategy(arb_two_hash_batches(100)))
    }

    fn fuzz(&self, data: &[u8]) {
        let input = fuzz_data_to_value(data, arb_two_hash_batches(100));
        test_proof_impl(input);
    }
}

#[derive(Clone, Debug, Default)]
pub struct AccumulatorConsistencyProof;

impl FuzzTargetImpl for AccumulatorConsistencyProof {
    fn description(&self) -> &'static str {
        "Accumulator consistency proof"
    }

    fn generate(&self, _idx: usize, _gen: &mut ValueGenerator) -> Option<Vec<u8>> {
        Some(corpus_from_strategy(arb_two_hash_batches(100)))
    }

    fn fuzz(&self, data: &[u8]) {
        let input = fuzz_data_to_value(data, arb_two_hash_batches(100));
        test_consistency_proof_impl(input);
    }
}

#[derive(Clone, Debug, Default)]
pub struct AccumulatorRangeProof;

impl FuzzTargetImpl for AccumulatorRangeProof {
    fn description(&self) -> &'static str {
        "Accumulator range proof"
    }

    fn generate(&self, _idx: usize, _gen: &mut ValueGenerator) -> Option<Vec<u8>> {
        Some(corpus_from_strategy(arb_three_hash_batches(100)))
    }

    fn fuzz(&self, data: &[u8]) {
        let input = fuzz_data_to_value(data, arb_three_hash_batches(100));
        test_range_proof_impl(input);
    }
}

#[derive(Clone, Debug, Default)]
pub struct AccumulatorAppendMany;

impl FuzzTargetImpl for AccumulatorAppendMany {
    fn description(&self) -> &'static str {
        "Accumulator amend many leaves"
    }

    fn generate(&self, _idx: usize, _gen: &mut ValueGenerator) -> Option<Vec<u8>> {
        Some(corpus_from_strategy(arb_list_of_hash_batches(10, 10)))
    }

    fn fuzz(&self, data: &[u8]) {
        let input = fuzz_data_to_value(data, arb_list_of_hash_batches(10, 10));
        test_append_many_impl(input);
    }
}

#[derive(Clone, Debug, Default)]
pub struct AccumulatorAppendEmpty;

impl FuzzTargetImpl for AccumulatorAppendEmpty {
    fn description(&self) -> &'static str {
        "Accumulator amend empty leave"
    }

    fn generate(&self, _idx: usize, _gen: &mut ValueGenerator) -> Option<Vec<u8>> {
        Some(corpus_from_strategy(arb_hash_batch(100)))
    }

    fn fuzz(&self, data: &[u8]) {
        let input = fuzz_data_to_value(data, arb_hash_batch(100));
        test_append_empty_impl(input);
    }
}
<|endoftext|>


<commit_before>#![cfg_attr(feature="serde_macros", feature(custom_derive, plugin))]
#![cfg_attr(feature="serde_macros", plugin(serde_macros))]

#[cfg(feature="serde_macros")]
include!("jsontypes.rs.in");

#[cfg(not(feature="serde_macros"))]
include!(concat!(env!("OUT_DIR"), "/jsontypes.rs"));
<commit_msg>Kill unused attributes<commit_after>#[cfg(feature="serde_macros")]
include!("jsontypes.rs.in");

#[cfg(not(feature="serde_macros"))]
include!(concat!(env!("OUT_DIR"), "/jsontypes.rs"));
<|endoftext|>


<commit_before>//! Process execution domains
use crate::Result;
use crate::errno::Errno;

use libc::{self, c_int, c_ulong};

libc_bitflags! {
    /// Flags used and returned by [`get()`](fn.get.html) and
    /// [`set()`](fn.set.html).
    pub struct Persona: c_int {
        /// Provide the legacy virtual address space layout.
        ADDR_COMPAT_LAYOUT;
        /// Disable address-space-layout randomization.
        ADDR_NO_RANDOMIZE;
        /// Limit the address space to 32 bits.
        ADDR_LIMIT_32BIT;
        /// Use `0xc0000000` as the offset at which to search a virtual memory
        /// chunk on [`mmap(2)`], otherwise use `0xffffe000`.
        ///
        /// [`mmap(2)`]: https://man7.org/linux/man-pages/man2/mmap.2.html
        ADDR_LIMIT_3GB;
        /// User-space function pointers to signal handlers point to descriptors.
        #[cfg(not(any(target_env = "musl", target_env = "uclibc")))]
        #[cfg_attr(docsrs, doc(cfg(all())))]
        FDPIC_FUNCPTRS;
        /// Map page 0 as read-only.
        MMAP_PAGE_ZERO;
        /// `PROT_READ` implies `PROT_EXEC` for [`mmap(2)`].
        ///
        /// [`mmap(2)`]: https://man7.org/linux/man-pages/man2/mmap.2.html
        READ_IMPLIES_EXEC;
        /// No effects.
        SHORT_INODE;
        /// [`select(2)`], [`pselect(2)`], and [`ppoll(2)`] do not modify the
        /// returned timeout argument when interrupted by a signal handler.
        ///
        /// [`select(2)`]: https://man7.org/linux/man-pages/man2/select.2.html
        /// [`pselect(2)`]: https://man7.org/linux/man-pages/man2/pselect.2.html
        /// [`ppoll(2)`]: https://man7.org/linux/man-pages/man2/ppoll.2.html
        STICKY_TIMEOUTS;
        /// Have [`uname(2)`] report a 2.6.40+ version number rather than a 3.x
        /// version number.
        ///
        /// [`uname(2)`]: https://man7.org/linux/man-pages/man2/uname.2.html
        #[cfg(not(any(target_env = "musl", target_env = "uclibc")))]
        #[cfg_attr(docsrs, doc(cfg(all())))]
        UNAME26;
        /// No effects.
        WHOLE_SECONDS;
    }
}

/// Retrieve the current process personality.
///
/// Returns a Result containing a Persona instance.
///
/// Example:
///
/// ```
/// # use nix::sys::personality::{self, Persona};
/// let pers = personality::get().unwrap();
/// assert!(!pers.contains(Persona::WHOLE_SECONDS));
/// ```
pub fn get() -> Result<Persona> {
    let res = unsafe {
        libc::personality(0xFFFFFFFF)
    };

    Errno::result(res).map(Persona::from_bits_truncate)
}

/// Set the current process personality.
///
/// Returns a Result containing the *previous* personality for the
/// process, as a Persona.
///
/// For more information, see [personality(2)](https://man7.org/linux/man-pages/man2/personality.2.html)
///
/// **NOTE**: This call **replaces** the current personality entirely.
/// To **update** the personality, first call `get()` and then `set()`
/// with the modified persona.
///
/// Example:
///
/// ```
/// # use nix::sys::personality::{self, Persona};
/// let mut pers = personality::get().unwrap();
/// assert!(!pers.contains(Persona::ADDR_NO_RANDOMIZE));
/// personality::set(pers | Persona::ADDR_NO_RANDOMIZE).unwrap().unwrap();
/// ```
pub fn set(persona: Persona) -> Result<Persona> {
    let res = unsafe {
        libc::personality(persona.bits() as c_ulong)
    };

    Errno::result(res).map(Persona::from_bits_truncate)
}
<commit_msg>fix CI error<commit_after>//! Process execution domains
use crate::Result;
use crate::errno::Errno;

use libc::{self, c_int, c_ulong};

libc_bitflags! {
    /// Flags used and returned by [`get()`](fn.get.html) and
    /// [`set()`](fn.set.html).
    pub struct Persona: c_int {
        /// Provide the legacy virtual address space layout.
        ADDR_COMPAT_LAYOUT;
        /// Disable address-space-layout randomization.
        ADDR_NO_RANDOMIZE;
        /// Limit the address space to 32 bits.
        ADDR_LIMIT_32BIT;
        /// Use `0xc0000000` as the offset at which to search a virtual memory
        /// chunk on [`mmap(2)`], otherwise use `0xffffe000`.
        ///
        /// [`mmap(2)`]: https://man7.org/linux/man-pages/man2/mmap.2.html
        ADDR_LIMIT_3GB;
        /// User-space function pointers to signal handlers point to descriptors.
        #[cfg(not(any(target_env = "musl", target_env = "uclibc")))]
        #[cfg_attr(docsrs, doc(cfg(all())))]
        FDPIC_FUNCPTRS;
        /// Map page 0 as read-only.
        MMAP_PAGE_ZERO;
        /// `PROT_READ` implies `PROT_EXEC` for [`mmap(2)`].
        ///
        /// [`mmap(2)`]: https://man7.org/linux/man-pages/man2/mmap.2.html
        READ_IMPLIES_EXEC;
        /// No effects.
        SHORT_INODE;
        /// [`select(2)`], [`pselect(2)`], and [`ppoll(2)`] do not modify the
        /// returned timeout argument when interrupted by a signal handler.
        ///
        /// [`select(2)`]: https://man7.org/linux/man-pages/man2/select.2.html
        /// [`pselect(2)`]: https://man7.org/linux/man-pages/man2/pselect.2.html
        /// [`ppoll(2)`]: https://man7.org/linux/man-pages/man2/ppoll.2.html
        STICKY_TIMEOUTS;
        /// Have [`uname(2)`] report a 2.6.40+ version number rather than a 3.x
        /// version number.
        ///
        /// [`uname(2)`]: https://man7.org/linux/man-pages/man2/uname.2.html
        #[cfg(not(any(target_env = "musl", target_env = "uclibc")))]
        #[cfg_attr(docsrs, doc(cfg(all())))]
        UNAME26;
        /// No effects.
        WHOLE_SECONDS;
    }
}

/// Retrieve the current process personality.
///
/// Returns a Result containing a Persona instance.
///
/// Example:
///
/// ```
/// # use nix::sys::personality::{self, Persona};
/// let pers = personality::get().unwrap();
/// assert!(!pers.contains(Persona::WHOLE_SECONDS));
/// ```
pub fn get() -> Result<Persona> {
    let res = unsafe {
        libc::personality(0xFFFFFFFF)
    };

    Errno::result(res).map(Persona::from_bits_truncate)
}

/// Set the current process personality.
///
/// Returns a Result containing the *previous* personality for the
/// process, as a Persona.
///
/// For more information, see [personality(2)](https://man7.org/linux/man-pages/man2/personality.2.html)
///
/// **NOTE**: This call **replaces** the current personality entirely.
/// To **update** the personality, first call `get()` and then `set()`
/// with the modified persona.
///
/// Example:
///
/// ```
/// # use nix::sys::personality::{self, Persona};
/// let mut pers = personality::get().unwrap();
/// assert!(!pers.contains(Persona::ADDR_NO_RANDOMIZE));
/// personality::set(pers | Persona::ADDR_NO_RANDOMIZE).unwrap();
/// ```
pub fn set(persona: Persona) -> Result<Persona> {
    let res = unsafe {
        libc::personality(persona.bits() as c_ulong)
    };

    Errno::result(res).map(Persona::from_bits_truncate)
}
<|endoftext|>


<commit_before>use std::os::unix::io::{AsRawFd, RawFd};
#[cfg(debug_assertions)]
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
use std::{cmp, io, mem, ptr};

use log::error;

use crate::sys::Events;
use crate::{Interests, Token};

/// Unique id for use as `SelectorId`.
#[cfg(debug_assertions)]
static NEXT_ID: AtomicUsize = AtomicUsize::new(1);

// Type of the `nchanges` and `nevents` parameters in the `kevent` function.
#[cfg(not(target_os = "netbsd"))]
type Count = libc::c_int;
#[cfg(target_os = "netbsd")]
type Count = libc::size_t;

// Type of the `filter` field in the `kevent` structure.
#[cfg(any(target_os = "freebsd", target_os = "openbsd"))]
type Filter = libc::c_short;
#[cfg(any(target_os = "macos", target_os = "ios"))]
type Filter = i16;
#[cfg(target_os = "netbsd")]
type Filter = u32;

// Type of the `data` field in the `kevent` structure.
#[cfg(any(
    target_os = "dragonfly",
    target_os = "freebsd",
    target_os = "ios",
    target_os = "macos"
))]
type Data = libc::intptr_t;
#[cfg(any(target_os = "bitrig", target_os = "netbsd", target_os = "openbsd",))]
type Data = i64;

// Type of the `udata` field in the `kevent` structure.
#[cfg(not(target_os = "netbsd"))]
type UData = *mut libc::c_void;
#[cfg(target_os = "netbsd")]
type UData = libc::intptr_t;

macro_rules! kevent {
    ($id: expr, $filter: expr, $flags: expr, $data: expr) => {
        libc::kevent {
            ident: $id as libc::uintptr_t,
            filter: $filter as Filter,
            flags: $flags,
            fflags: 0,
            data: 0,
            udata: $data as UData,
        }
    };
}

#[derive(Debug)]
pub struct Selector {
    #[cfg(debug_assertions)]
    id: usize,
    kq: RawFd,
}

impl Selector {
    pub fn new() -> io::Result<Selector> {
        syscall!(kqueue())
            .and_then(|kq| syscall!(fcntl(kq, libc::F_SETFD, libc::FD_CLOEXEC)).map(|_| kq))
            .map(|kq| Selector {
                #[cfg(debug_assertions)]
                id: NEXT_ID.fetch_add(1, Ordering::Relaxed),
                kq,
            })
    }

    #[cfg(debug_assertions)]
    pub fn id(&self) -> usize {
        self.id
    }

    pub fn select(&self, events: &mut Events, timeout: Option<Duration>) -> io::Result<()> {
        let timeout = timeout.map(|to| libc::timespec {
            tv_sec: cmp::min(to.as_secs(), libc::time_t::max_value() as u64) as libc::time_t,
            // `Duration::subsec_nanos` is guaranteed to be less than one
            // billion (the number of nanoseconds in a second), making the
            // cast to i32 safe. The cast itself is needed for platforms
            // where C's long is only 32 bits.
            tv_nsec: libc::c_long::from(to.subsec_nanos() as i32),
        });
        let timeout = timeout
            .as_ref()
            .map(|s| s as *const _)
            .unwrap_or(ptr::null_mut());

        events.clear();
        syscall!(kevent(
            self.kq,
            ptr::null(),
            0,
            events.as_mut_ptr(),
            events.capacity() as Count,
            timeout,
        ))
        .map(|n_events| {
            // This is safe because `kevent` ensures that `n_events` are
            // assigned.
            unsafe { events.set_len(n_events as usize) };
        })
    }

    pub fn register(&self, fd: RawFd, token: Token, interests: Interests) -> io::Result<()> {
        let flags = libc::EV_CLEAR | libc::EV_RECEIPT | libc::EV_ADD;
        // At most we need two changes, but maybe we only need 1.
        // FIXME: replace this with `MaybeUninit` once that become usable with
        // arrays.
        let mut changes: [libc::kevent; 2] = unsafe { mem::uninitialized() };
        let mut n_changes = 0;

        if interests.is_writable() {
            let kevent = kevent!(fd, libc::EVFILT_WRITE, flags, token.0);
            unsafe { ptr::write(&mut changes[n_changes], kevent) };
            n_changes += 1;
        }

        if interests.is_readable() {
            let kevent = kevent!(fd, libc::EVFILT_READ, flags, token.0);
            unsafe { ptr::write(&mut changes[n_changes], kevent) };
            n_changes += 1;
        }

        // Older versions of macOS (OS X 10.11 and 10.10 have been witnessed)
        // can return EPIPE when registering a pipe file descriptor where the
        // other end has already disappeared. For example code that creates a
        // pipe, closes a file descriptor, and then registers the other end will
        // see an EPIPE returned from `register`.
        //
        // It also turns out that kevent will still report events on the file
        // descriptor, telling us that it's readable/hup at least after we've
        // done this registration. As a result we just ignore `EPIPE` here
        // instead of propagating it.
        //
        // More info can be found at tokio-rs/mio#582.
        kevent_register(self.kq, &mut changes[0..n_changes], &[libc::EPIPE as Data])
    }

    pub fn reregister(&self, fd: RawFd, token: Token, interests: Interests) -> io::Result<()> {
        let flags = libc::EV_CLEAR | libc::EV_RECEIPT;
        let write_flags = if interests.is_writable() {
            flags | libc::EV_ADD
        } else {
            flags | libc::EV_DELETE
        };
        let read_flags = if interests.is_readable() {
            flags | libc::EV_ADD
        } else {
            flags | libc::EV_DELETE
        };

        let mut changes: [libc::kevent; 2] = [
            kevent!(fd, libc::EVFILT_WRITE, write_flags, token.0),
            kevent!(fd, libc::EVFILT_READ, read_flags, token.0),
        ];

        // Since there is no way to check with which interests the fd was
        // registered we modify both readable and write, adding it when required
        // and removing it otherwise, ignoring the ENOENT error when it comes
        // up. The ENOENT error informs us that a filter we're trying to remove
        // wasn't there in first place, but we don't really care since our goal
        // is accomplished.
        //
        // For the explanation of ignoring `EPIPE` see `register`.
        kevent_register(
            self.kq,
            &mut changes,
            &[libc::ENOENT as Data, libc::EPIPE as Data],
        )
    }

    pub fn deregister(&self, fd: RawFd) -> io::Result<()> {
        let flags = libc::EV_DELETE | libc::EV_RECEIPT;
        let mut changes: [libc::kevent; 2] = [
            kevent!(fd, libc::EVFILT_WRITE, flags, 0),
            kevent!(fd, libc::EVFILT_READ, flags, 0),
        ];

        // Since there is no way to check with which interests the fd was
        // registered we remove both filters (readable and writeable) and ignore
        // the ENOENT error when it comes up. The ENOENT error informs us that
        // the filter wasn't there in first place, but we don't really care
        // about that since our goal is to remove it.
        kevent_register(self.kq, &mut changes, &[libc::ENOENT as Data])
    }

    // Used by `Waker`.
    #[cfg(any(target_os = "freebsd", target_os = "ios", target_os = "macos"))]
    pub fn setup_waker(&self, token: Token) -> io::Result<()> {
        // First attempt to accept user space notifications.
        let mut kevent = kevent!(
            0,
            libc::EVFILT_USER,
            libc::EV_ADD | libc::EV_CLEAR | libc::EV_RECEIPT,
            token.0
        );

        syscall!(kevent(self.kq, &kevent, 1, &mut kevent, 1, ptr::null())).and_then(|_| {
            if (kevent.flags & libc::EV_ERROR) != 0 && kevent.data != 0 {
                Err(io::Error::from_raw_os_error(kevent.data as i32))
            } else {
                Ok(())
            }
        })
    }

    // Used by `Waker`.
    #[cfg(any(target_os = "freebsd", target_os = "ios", target_os = "macos"))]
    pub fn try_clone_waker(&self) -> io::Result<Selector> {
        syscall!(dup(self.kq)).map(|new_kq| Selector {
            #[cfg(debug_assertions)]
            id: self.id,
            kq: new_kq,
        })
    }

    // Used by `Waker`.
    #[cfg(any(target_os = "freebsd", target_os = "ios", target_os = "macos"))]
    pub fn wake(&self, token: Token) -> io::Result<()> {
        let mut kevent = kevent!(
            0,
            libc::EVFILT_USER,
            libc::EV_ADD | libc::EV_RECEIPT,
            token.0
        );
        kevent.fflags = libc::NOTE_TRIGGER;

        syscall!(kevent(self.kq, &kevent, 1, &mut kevent, 1, ptr::null())).and_then(|_| {
            if (kevent.flags & libc::EV_ERROR) != 0 && kevent.data != 0 {
                Err(io::Error::from_raw_os_error(kevent.data as i32))
            } else {
                Ok(())
            }
        })
    }
}

/// Register `changes` with `kq`ueue.
fn kevent_register(
    kq: RawFd,
    changes: &mut [libc::kevent],
    ignored_errors: &[Data],
) -> io::Result<()> {
    syscall!(kevent(
        kq,
        changes.as_ptr(),
        changes.len() as Count,
        changes.as_mut_ptr(),
        changes.len() as Count,
        ptr::null(),
    ))
    .map(|_| ())
    .or_else(|err| {
        // According to the manual page of FreeBSD: "When kevent() call fails
        // with EINTR error, all changes in the changelist have been applied",
        // so we can safely ignore it.
        if err.raw_os_error() == Some(libc::EINTR) {
            Ok(())
        } else {
            Err(err)
        }
    })
    .and_then(|()| check_errors(&changes, ignored_errors))
}

/// Check all events for possible errors, it returns the first error found.
fn check_errors(events: &[libc::kevent], ignored_errors: &[Data]) -> io::Result<()> {
    for event in events {
        // We can't use references to packed structures (in checking the ignored
        // errors), so we need copy the data out before use.
        let data = event.data;
        // Check for the error flag, the actual error will be in the `data`
        // field.
        if (event.flags & libc::EV_ERROR != 0) && data != 0 && !ignored_errors.contains(&data) {
            return Err(io::Error::from_raw_os_error(data as i32));
        }
    }
    Ok(())
}

impl AsRawFd for Selector {
    fn as_raw_fd(&self) -> RawFd {
        self.kq
    }
}

impl Drop for Selector {
    fn drop(&mut self) {
        if let Err(err) = syscall!(close(self.kq)) {
            error!("error closing kqueue: {}", err);
        }
    }
}

pub type Event = libc::kevent;

pub mod event {
    use crate::sys::Event;
    use crate::Token;

    pub fn token(event: &Event) -> Token {
        Token(event.udata as usize)
    }

    pub fn is_readable(event: &Event) -> bool {
        event.filter == libc::EVFILT_READ || {
            #[cfg(any(target_os = "freebsd", target_os = "ios", target_os = "macos"))]
            // Used by the `Awakener`. On platforms that use `eventfd` or a unix
            // pipe it will emit a readable event so we'll fake that here as
            // well.
            {
                event.filter == libc::EVFILT_USER
            }
            #[cfg(not(any(target_os = "freebsd", target_os = "ios", target_os = "macos")))]
            {
                false
            }
        }
    }

    pub fn is_writable(event: &Event) -> bool {
        event.filter == libc::EVFILT_WRITE
    }

    pub fn is_error(event: &Event) -> bool {
        (event.flags & libc::EV_ERROR) != 0 ||
            // When the read end of the socket is closed, EV_EOF is set on
            // flags, and fflags contains the error if there is one.
            (event.flags & libc::EV_EOF) != 0 && event.fflags != 0
    }

    pub fn is_hup(event: &Event) -> bool {
        (event.flags & libc::EV_EOF) != 0
    }

    pub fn is_priority(_: &Event) -> bool {
        // kqueue doesn't have priority indicators.
        false
    }

    #[allow(unused_variables)] // `event` is not used on some platforms.
    pub fn is_aio(event: &Event) -> bool {
        #[cfg(any(
            target_os = "dragonfly",
            target_os = "freebsd",
            target_os = "ios",
            target_os = "macos"
        ))]
        {
            event.filter == libc::EVFILT_AIO
        }
        #[cfg(not(any(
            target_os = "dragonfly",
            target_os = "freebsd",
            target_os = "ios",
            target_os = "macos"
        )))]
        {
            false
        }
    }

    #[allow(unused_variables)] // `event` is only used on FreeBSD.
    pub fn is_lio(event: &Event) -> bool {
        #[cfg(target_os = "freebsd")]
        {
            event.filter == libc::EVFILT_LIO
        }
        #[cfg(not(target_os = "freebsd"))]
        {
            false
        }
    }
}

#[test]
fn does_not_register_rw() {
    use crate::unix::SourceFd;
    use crate::{Poll, Token};

    let kq = unsafe { libc::kqueue() };
    let kqf = SourceFd(&kq);
    let poll = Poll::new().unwrap();

    // Registering kqueue fd will fail if write is requested (On anything but
    // some versions of macOS).
    poll.registry()
        .register(&kqf, Token(1234), Interests::READABLE)
        .unwrap();
}
<commit_msg>Replace uninitialized with MaybeUninit<commit_after>use std::mem::MaybeUninit;
use std::os::unix::io::{AsRawFd, RawFd};
#[cfg(debug_assertions)]
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
use std::{cmp, io, ptr, slice};

use log::error;

use crate::sys::Events;
use crate::{Interests, Token};

/// Unique id for use as `SelectorId`.
#[cfg(debug_assertions)]
static NEXT_ID: AtomicUsize = AtomicUsize::new(1);

// Type of the `nchanges` and `nevents` parameters in the `kevent` function.
#[cfg(not(target_os = "netbsd"))]
type Count = libc::c_int;
#[cfg(target_os = "netbsd")]
type Count = libc::size_t;

// Type of the `filter` field in the `kevent` structure.
#[cfg(any(target_os = "freebsd", target_os = "openbsd"))]
type Filter = libc::c_short;
#[cfg(any(target_os = "macos", target_os = "ios"))]
type Filter = i16;
#[cfg(target_os = "netbsd")]
type Filter = u32;

// Type of the `data` field in the `kevent` structure.
#[cfg(any(
    target_os = "dragonfly",
    target_os = "freebsd",
    target_os = "ios",
    target_os = "macos"
))]
type Data = libc::intptr_t;
#[cfg(any(target_os = "bitrig", target_os = "netbsd", target_os = "openbsd",))]
type Data = i64;

// Type of the `udata` field in the `kevent` structure.
#[cfg(not(target_os = "netbsd"))]
type UData = *mut libc::c_void;
#[cfg(target_os = "netbsd")]
type UData = libc::intptr_t;

macro_rules! kevent {
    ($id: expr, $filter: expr, $flags: expr, $data: expr) => {
        libc::kevent {
            ident: $id as libc::uintptr_t,
            filter: $filter as Filter,
            flags: $flags,
            fflags: 0,
            data: 0,
            udata: $data as UData,
        }
    };
}

#[derive(Debug)]
pub struct Selector {
    #[cfg(debug_assertions)]
    id: usize,
    kq: RawFd,
}

impl Selector {
    pub fn new() -> io::Result<Selector> {
        syscall!(kqueue())
            .and_then(|kq| syscall!(fcntl(kq, libc::F_SETFD, libc::FD_CLOEXEC)).map(|_| kq))
            .map(|kq| Selector {
                #[cfg(debug_assertions)]
                id: NEXT_ID.fetch_add(1, Ordering::Relaxed),
                kq,
            })
    }

    #[cfg(debug_assertions)]
    pub fn id(&self) -> usize {
        self.id
    }

    pub fn select(&self, events: &mut Events, timeout: Option<Duration>) -> io::Result<()> {
        let timeout = timeout.map(|to| libc::timespec {
            tv_sec: cmp::min(to.as_secs(), libc::time_t::max_value() as u64) as libc::time_t,
            // `Duration::subsec_nanos` is guaranteed to be less than one
            // billion (the number of nanoseconds in a second), making the
            // cast to i32 safe. The cast itself is needed for platforms
            // where C's long is only 32 bits.
            tv_nsec: libc::c_long::from(to.subsec_nanos() as i32),
        });
        let timeout = timeout
            .as_ref()
            .map(|s| s as *const _)
            .unwrap_or(ptr::null_mut());

        events.clear();
        syscall!(kevent(
            self.kq,
            ptr::null(),
            0,
            events.as_mut_ptr(),
            events.capacity() as Count,
            timeout,
        ))
        .map(|n_events| {
            // This is safe because `kevent` ensures that `n_events` are
            // assigned.
            unsafe { events.set_len(n_events as usize) };
        })
    }

    pub fn register(&self, fd: RawFd, token: Token, interests: Interests) -> io::Result<()> {
        let flags = libc::EV_CLEAR | libc::EV_RECEIPT | libc::EV_ADD;
        // At most we need two changes, but maybe we only need 1.
        let mut changes: [MaybeUninit<libc::kevent>; 2] =
            [MaybeUninit::uninit(), MaybeUninit::uninit()];
        let mut n_changes = 0;

        if interests.is_writable() {
            let kevent = kevent!(fd, libc::EVFILT_WRITE, flags, token.0);
            changes[n_changes] = MaybeUninit::new(kevent);
            n_changes += 1;
        }

        if interests.is_readable() {
            let kevent = kevent!(fd, libc::EVFILT_READ, flags, token.0);
            changes[n_changes] = MaybeUninit::new(kevent);
            n_changes += 1;
        }

        // Older versions of macOS (OS X 10.11 and 10.10 have been witnessed)
        // can return EPIPE when registering a pipe file descriptor where the
        // other end has already disappeared. For example code that creates a
        // pipe, closes a file descriptor, and then registers the other end will
        // see an EPIPE returned from `register`.
        //
        // It also turns out that kevent will still report events on the file
        // descriptor, telling us that it's readable/hup at least after we've
        // done this registration. As a result we just ignore `EPIPE` here
        // instead of propagating it.
        //
        // More info can be found at tokio-rs/mio#582.
        let changes = unsafe {
            // This is safe because we ensure that at least `n_changes` are in
            // the array.
            slice::from_raw_parts_mut(changes[0].as_mut_ptr(), n_changes)
        };
        kevent_register(self.kq, changes, &[libc::EPIPE as Data])
    }

    pub fn reregister(&self, fd: RawFd, token: Token, interests: Interests) -> io::Result<()> {
        let flags = libc::EV_CLEAR | libc::EV_RECEIPT;
        let write_flags = if interests.is_writable() {
            flags | libc::EV_ADD
        } else {
            flags | libc::EV_DELETE
        };
        let read_flags = if interests.is_readable() {
            flags | libc::EV_ADD
        } else {
            flags | libc::EV_DELETE
        };

        let mut changes: [libc::kevent; 2] = [
            kevent!(fd, libc::EVFILT_WRITE, write_flags, token.0),
            kevent!(fd, libc::EVFILT_READ, read_flags, token.0),
        ];

        // Since there is no way to check with which interests the fd was
        // registered we modify both readable and write, adding it when required
        // and removing it otherwise, ignoring the ENOENT error when it comes
        // up. The ENOENT error informs us that a filter we're trying to remove
        // wasn't there in first place, but we don't really care since our goal
        // is accomplished.
        //
        // For the explanation of ignoring `EPIPE` see `register`.
        kevent_register(
            self.kq,
            &mut changes,
            &[libc::ENOENT as Data, libc::EPIPE as Data],
        )
    }

    pub fn deregister(&self, fd: RawFd) -> io::Result<()> {
        let flags = libc::EV_DELETE | libc::EV_RECEIPT;
        let mut changes: [libc::kevent; 2] = [
            kevent!(fd, libc::EVFILT_WRITE, flags, 0),
            kevent!(fd, libc::EVFILT_READ, flags, 0),
        ];

        // Since there is no way to check with which interests the fd was
        // registered we remove both filters (readable and writeable) and ignore
        // the ENOENT error when it comes up. The ENOENT error informs us that
        // the filter wasn't there in first place, but we don't really care
        // about that since our goal is to remove it.
        kevent_register(self.kq, &mut changes, &[libc::ENOENT as Data])
    }

    // Used by `Waker`.
    #[cfg(any(target_os = "freebsd", target_os = "ios", target_os = "macos"))]
    pub fn setup_waker(&self, token: Token) -> io::Result<()> {
        // First attempt to accept user space notifications.
        let mut kevent = kevent!(
            0,
            libc::EVFILT_USER,
            libc::EV_ADD | libc::EV_CLEAR | libc::EV_RECEIPT,
            token.0
        );

        syscall!(kevent(self.kq, &kevent, 1, &mut kevent, 1, ptr::null())).and_then(|_| {
            if (kevent.flags & libc::EV_ERROR) != 0 && kevent.data != 0 {
                Err(io::Error::from_raw_os_error(kevent.data as i32))
            } else {
                Ok(())
            }
        })
    }

    // Used by `Waker`.
    #[cfg(any(target_os = "freebsd", target_os = "ios", target_os = "macos"))]
    pub fn try_clone_waker(&self) -> io::Result<Selector> {
        syscall!(dup(self.kq)).map(|new_kq| Selector {
            #[cfg(debug_assertions)]
            id: self.id,
            kq: new_kq,
        })
    }

    // Used by `Waker`.
    #[cfg(any(target_os = "freebsd", target_os = "ios", target_os = "macos"))]
    pub fn wake(&self, token: Token) -> io::Result<()> {
        let mut kevent = kevent!(
            0,
            libc::EVFILT_USER,
            libc::EV_ADD | libc::EV_RECEIPT,
            token.0
        );
        kevent.fflags = libc::NOTE_TRIGGER;

        syscall!(kevent(self.kq, &kevent, 1, &mut kevent, 1, ptr::null())).and_then(|_| {
            if (kevent.flags & libc::EV_ERROR) != 0 && kevent.data != 0 {
                Err(io::Error::from_raw_os_error(kevent.data as i32))
            } else {
                Ok(())
            }
        })
    }
}

/// Register `changes` with `kq`ueue.
fn kevent_register(
    kq: RawFd,
    changes: &mut [libc::kevent],
    ignored_errors: &[Data],
) -> io::Result<()> {
    syscall!(kevent(
        kq,
        changes.as_ptr(),
        changes.len() as Count,
        changes.as_mut_ptr(),
        changes.len() as Count,
        ptr::null(),
    ))
    .map(|_| ())
    .or_else(|err| {
        // According to the manual page of FreeBSD: "When kevent() call fails
        // with EINTR error, all changes in the changelist have been applied",
        // so we can safely ignore it.
        if err.raw_os_error() == Some(libc::EINTR) {
            Ok(())
        } else {
            Err(err)
        }
    })
    .and_then(|()| check_errors(&changes, ignored_errors))
}

/// Check all events for possible errors, it returns the first error found.
fn check_errors(events: &[libc::kevent], ignored_errors: &[Data]) -> io::Result<()> {
    for event in events {
        // We can't use references to packed structures (in checking the ignored
        // errors), so we need copy the data out before use.
        let data = event.data;
        // Check for the error flag, the actual error will be in the `data`
        // field.
        if (event.flags & libc::EV_ERROR != 0) && data != 0 && !ignored_errors.contains(&data) {
            return Err(io::Error::from_raw_os_error(data as i32));
        }
    }
    Ok(())
}

impl AsRawFd for Selector {
    fn as_raw_fd(&self) -> RawFd {
        self.kq
    }
}

impl Drop for Selector {
    fn drop(&mut self) {
        if let Err(err) = syscall!(close(self.kq)) {
            error!("error closing kqueue: {}", err);
        }
    }
}

pub type Event = libc::kevent;

pub mod event {
    use crate::sys::Event;
    use crate::Token;

    pub fn token(event: &Event) -> Token {
        Token(event.udata as usize)
    }

    pub fn is_readable(event: &Event) -> bool {
        event.filter == libc::EVFILT_READ || {
            #[cfg(any(target_os = "freebsd", target_os = "ios", target_os = "macos"))]
            // Used by the `Awakener`. On platforms that use `eventfd` or a unix
            // pipe it will emit a readable event so we'll fake that here as
            // well.
            {
                event.filter == libc::EVFILT_USER
            }
            #[cfg(not(any(target_os = "freebsd", target_os = "ios", target_os = "macos")))]
            {
                false
            }
        }
    }

    pub fn is_writable(event: &Event) -> bool {
        event.filter == libc::EVFILT_WRITE
    }

    pub fn is_error(event: &Event) -> bool {
        (event.flags & libc::EV_ERROR) != 0 ||
            // When the read end of the socket is closed, EV_EOF is set on
            // flags, and fflags contains the error if there is one.
            (event.flags & libc::EV_EOF) != 0 && event.fflags != 0
    }

    pub fn is_hup(event: &Event) -> bool {
        (event.flags & libc::EV_EOF) != 0
    }

    pub fn is_priority(_: &Event) -> bool {
        // kqueue doesn't have priority indicators.
        false
    }

    #[allow(unused_variables)] // `event` is not used on some platforms.
    pub fn is_aio(event: &Event) -> bool {
        #[cfg(any(
            target_os = "dragonfly",
            target_os = "freebsd",
            target_os = "ios",
            target_os = "macos"
        ))]
        {
            event.filter == libc::EVFILT_AIO
        }
        #[cfg(not(any(
            target_os = "dragonfly",
            target_os = "freebsd",
            target_os = "ios",
            target_os = "macos"
        )))]
        {
            false
        }
    }

    #[allow(unused_variables)] // `event` is only used on FreeBSD.
    pub fn is_lio(event: &Event) -> bool {
        #[cfg(target_os = "freebsd")]
        {
            event.filter == libc::EVFILT_LIO
        }
        #[cfg(not(target_os = "freebsd"))]
        {
            false
        }
    }
}

#[test]
fn does_not_register_rw() {
    use crate::unix::SourceFd;
    use crate::{Poll, Token};

    let kq = unsafe { libc::kqueue() };
    let kqf = SourceFd(&kq);
    let poll = Poll::new().unwrap();

    // Registering kqueue fd will fail if write is requested (On anything but
    // some versions of macOS).
    poll.registry()
        .register(&kqf, Token(1234), Interests::READABLE)
        .unwrap();
}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2016 Sandstorm Development Group, Inc. and contributors
// Licensed under the MIT License:
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

use gj::{Promise, PromiseFulfiller};
use std::cell::{RefCell};
use std::rc::Rc;
use handle_table::{HandleTable, Handle};

struct Observer {
    read_overlapped: *mut ::miow::Overlapped,
    write_overlapped: *mut ::miow::Overlapped,
    read_fulfiller: Option<PromiseFulfiller<u32, ::std::io::Error>>,
    write_fulfiller: Option<PromiseFulfiller<u32, ::std::io::Error>>,
}

impl Observer {
    pub fn when_read_done(&mut self) -> Promise<u32, ::std::io::Error> {
        let (promise, fulfiller) = Promise::and_fulfiller();
        self.read_fulfiller = Some(fulfiller);
        promise
    }

    pub fn when_write_done(&mut self) -> Promise<u32, ::std::io::Error> {
        let (promise, fulfiller) = Promise::and_fulfiller();
        self.write_fulfiller = Some(fulfiller);
        promise
    }
}

pub struct Reactor {
    cp: ::miow::iocp::CompletionPort,
    observers: HandleTable<Observer>,
    statuses: Vec<::miow::iocp::CompletionStatus>,
}

impl Reactor {
    pub fn new() -> Result<Reactor, ::std::io::Error> {
        Ok(Reactor {
            cp: try!(::miow::iocp::CompletionPort::new(1)),
            observers: HandleTable::new(),
            statuses: vec![::miow::iocp::CompletionStatus::zero(); 1024]
        })
    }

    pub fn run_once(&mut self, maybe_timeout: Option<::time::Duration>)
                    -> Result<(), ::std::io::Error>
    {
        let timeout = maybe_timeout.map(|t| t.num_milliseconds() as u32); // XXX check for overflow

        //let mut statu

        {
            let statuses = try!(self.cp.get_many(&mut self.statuses[..], timeout));
            for status in statuses {
                let token = status.token();
                let overlapped = status.overlapped();
                let bytes_transferred = status.bytes_transferred();
                let handle = Handle { val: token };
                if self.observers[handle].read_overlapped == overlapped {
                    match self.observers[handle].read_fulfiller.take() {
                        None => (),
                        Some(f) => f.fulfill(bytes_transferred),
                    }
                }
                if self.observers[handle].write_overlapped == overlapped {
                    match self.observers[handle].write_fulfiller.take() {
                        None => (),
                        Some(f) => f.fulfill(bytes_transferred),
                    }
                }

            }
        }
        Ok(())
    }

    fn add_socket<T>(&mut self, sock: &T,
                     read_overlapped: *mut ::miow::Overlapped,
                     write_overlapped: *mut ::miow::Overlapped)
                     -> Result<Handle, ::std::io::Error>
        where T : ::std::os::windows::io::AsRawSocket
    {
        let observer = Observer {
            read_overlapped: read_overlapped,
            write_overlapped: write_overlapped,
            read_fulfiller: None,
            write_fulfiller: None,
        };
        let handle = self.observers.push(observer);
        try!(self.cp.add_socket(handle.val, sock));
        Ok(handle)
    }
}

#[derive(Clone)]
pub struct SocketAddressInner {
    reactor: Rc<RefCell<Reactor>>,
    addr: ::std::net::SocketAddr,
}

impl SocketAddressInner {
    pub fn new_tcp(reactor: Rc<RefCell<Reactor>>, addr: ::std::net::SocketAddr)
                   -> SocketAddressInner
    {
        SocketAddressInner {
            reactor: reactor,
            addr: addr,
        }
    }

    pub fn connect(&self) -> Promise<SocketStreamInner, ::std::io::Error> {
        use miow::net::TcpBuilderExt;
        let builder = match self.addr {
            ::std::net::SocketAddr::V4(_) => pry!(::net2::TcpBuilder::new_v4()),
            ::std::net::SocketAddr::V6(_) => pry!(::net2::TcpBuilder::new_v6()),
        };

        let mut read_overlapped = Box::new(::miow::Overlapped::zero());
        let mut write_overlapped = Box::new(::miow::Overlapped::zero());
        let handle = pry!(
            self.reactor.borrow_mut().add_socket(&builder, &mut *read_overlapped,
                                                 &mut *write_overlapped));

        let (stream, _ready) = unsafe {
            pry!(builder.connect_overlapped(&self.addr, &mut *write_overlapped))
        };

        let reactor2 = self.reactor.clone();
        let result = self.reactor.borrow_mut().observers[handle].when_write_done().map(move |_| {
            Ok(SocketStreamInner::new_already_registered(
                reactor2, stream,
                read_overlapped, write_overlapped,
                handle))
        });

        result
    }

    pub fn listen(&mut self) -> Result<SocketListenerInner, ::std::io::Error> {
        let listener = try!(::std::net::TcpListener::bind(self.addr));
        SocketListenerInner::new(self.reactor.clone(), listener, self.addr)
    }
}

pub struct SocketListenerInner {
    reactor: Rc<RefCell<Reactor>>,
    listener: ::std::net::TcpListener,
    addr: ::std::net::SocketAddr,
    read_overlapped: Box<::miow::Overlapped>,
    handle: Handle,
    pub queue: Option<Promise<(),()>>,
}

impl SocketListenerInner {
    fn new(reactor: Rc<RefCell<Reactor>>, listener: ::std::net::TcpListener,
           addr: ::std::net::SocketAddr)
           -> Result<SocketListenerInner, ::std::io::Error>
    {
        let mut read_overlapped = Box::new(::miow::Overlapped::zero());
        let handle = try!(
            reactor.borrow_mut().add_socket(&listener, &mut *read_overlapped,
                                            ::std::ptr::null_mut()));

        Ok(SocketListenerInner {
            reactor: reactor,
            listener: listener,
            addr: addr,
            read_overlapped: read_overlapped,
            handle: handle,
            queue: None,
        })
    }

    pub fn accept_internal(inner: Rc<RefCell<SocketListenerInner>>)
                           -> Promise<SocketStreamInner, ::std::io::Error>
    {
        use miow::net::TcpListenerExt;
        let builder = match inner.borrow().addr {
            ::std::net::SocketAddr::V4(_) => pry!(::net2::TcpBuilder::new_v4()),
            ::std::net::SocketAddr::V6(_) => pry!(::net2::TcpBuilder::new_v6()),
        };

        let mut accept_addrs = ::miow::net::AcceptAddrsBuf::new();

        let &mut SocketListenerInner {
            ref reactor, ref mut listener, ref mut read_overlapped, handle,
            ..
        } =  &mut *inner.borrow_mut();

        let (stream, _ready) = unsafe {
            pry!(listener.accept_overlapped(&builder,
                                            &mut accept_addrs,
                                            read_overlapped))
        };

        let reactor2 = reactor.clone();
        let result = reactor.borrow_mut().observers[handle].when_read_done().map(move |_| {
            println!("accepted!");
            SocketStreamInner::new(reactor2, stream)
        });

        result
    }
}

pub struct SocketStreamInner {
    reactor: Rc<RefCell<Reactor>>,
    stream: ::std::net::TcpStream,
    read_overlapped: Box<::miow::Overlapped>,
    write_overlapped: Box<::miow::Overlapped>,
    handle: Handle,
    pub read_queue: Option<Promise<(),()>>,
    pub write_queue: Option<Promise<(),()>>,
}

impl SocketStreamInner {
    fn new(reactor: Rc<RefCell<Reactor>>, stream: ::std::net::TcpStream)
           -> Result<SocketStreamInner, ::std::io::Error>
    {
        let mut read_overlapped = Box::new(::miow::Overlapped::zero());
        let mut write_overlapped = Box::new(::miow::Overlapped::zero());
        let handle = try!(
            reactor.borrow_mut().add_socket(&stream, &mut *read_overlapped,
                                            &mut *write_overlapped));
        Ok(SocketStreamInner {
            reactor: reactor,
            stream: stream,
            read_overlapped: read_overlapped,
            write_overlapped: write_overlapped,
            handle: handle,
            read_queue: None,
            write_queue: None,
        })
    }

    fn new_already_registered(reactor: Rc<RefCell<Reactor>>, stream: ::std::net::TcpStream,
                              read_overlapped: Box<::miow::Overlapped>,
                              write_overlapped: Box<::miow::Overlapped>,
                              handle: Handle)
           -> SocketStreamInner
    {
        SocketStreamInner {
            reactor: reactor,
            stream: stream,
            read_overlapped: read_overlapped,
            write_overlapped: write_overlapped,
            handle: handle,
            read_queue: None,
            write_queue: None,
        }
    }

    pub fn try_read_internal<T>(inner: Rc<RefCell<SocketStreamInner>>,
                                mut buf: T,
                                already_read: usize,
                                min_bytes: usize)
                                -> Promise<(T, usize), ::std::io::Error>
        where T: AsMut<[u8]>
    {
        use ::miow::net::TcpStreamExt;

        if already_read >= min_bytes {
            return Promise::ok((buf, already_read));
        }

        let inner2 = inner.clone();
        let &mut SocketStreamInner {
            ref reactor, ref mut stream,
            ref mut read_overlapped, handle,
            ..
        } =  &mut *inner.borrow_mut();

        let _done = unsafe {
            pry!(stream.read_overlapped(&mut buf.as_mut()[already_read..], read_overlapped))
        };

        let result = reactor.borrow_mut().observers[handle].when_read_done().then(move |n| {
            println!("read transferred this many bytes: {}", n);
            let total_read = n as usize + already_read;
            if n == 0 {
                println!("EOF");
                Promise::ok((buf, total_read))
            } else {
                SocketStreamInner::try_read_internal(inner2, buf, total_read, min_bytes)
            }
        });

        result
    }

    pub fn write_internal<T>(inner: Rc<RefCell<SocketStreamInner>>,
                             buf: T,
                             already_written: usize) -> Promise<T, ::std::io::Error>
        where T: AsRef<[u8]>
    {
        use ::miow::net::TcpStreamExt;

        if already_written == buf.as_ref().len() {
            return Promise::ok(buf)
        }

        let inner2 = inner.clone();
        let &mut SocketStreamInner {
            ref reactor,
            ref stream,
            ref mut write_overlapped, handle,
            ..
        } =  &mut *inner.borrow_mut();

        let _done = unsafe {
            pry!(stream.write_overlapped(&buf.as_ref()[already_written ..], write_overlapped))
        };

        let result = reactor.borrow_mut().observers[handle].when_write_done().then(move |n| {
            println!("write transferred this many bytes: {}", n);
            let total_written = n as usize + already_written;
            if n == 0 {
                println!("wrote zero bytes!?");
            }

            SocketStreamInner::write_internal(inner2, buf, total_written)
        });

        result
    }
}
<commit_msg>bind the socket<commit_after>// Copyright (c) 2013-2016 Sandstorm Development Group, Inc. and contributors
// Licensed under the MIT License:
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

use gj::{Promise, PromiseFulfiller};
use std::cell::{RefCell};
use std::rc::Rc;
use handle_table::{HandleTable, Handle};

struct Observer {
    read_overlapped: *mut ::miow::Overlapped,
    write_overlapped: *mut ::miow::Overlapped,
    read_fulfiller: Option<PromiseFulfiller<u32, ::std::io::Error>>,
    write_fulfiller: Option<PromiseFulfiller<u32, ::std::io::Error>>,
}

impl Observer {
    pub fn when_read_done(&mut self) -> Promise<u32, ::std::io::Error> {
        let (promise, fulfiller) = Promise::and_fulfiller();
        self.read_fulfiller = Some(fulfiller);
        promise
    }

    pub fn when_write_done(&mut self) -> Promise<u32, ::std::io::Error> {
        let (promise, fulfiller) = Promise::and_fulfiller();
        self.write_fulfiller = Some(fulfiller);
        promise
    }
}

pub struct Reactor {
    cp: ::miow::iocp::CompletionPort,
    observers: HandleTable<Observer>,
    statuses: Vec<::miow::iocp::CompletionStatus>,
}

impl Reactor {
    pub fn new() -> Result<Reactor, ::std::io::Error> {
        Ok(Reactor {
            cp: try!(::miow::iocp::CompletionPort::new(1)),
            observers: HandleTable::new(),
            statuses: vec![::miow::iocp::CompletionStatus::zero(); 1024]
        })
    }

    pub fn run_once(&mut self, maybe_timeout: Option<::time::Duration>)
                    -> Result<(), ::std::io::Error>
    {
        let timeout = maybe_timeout.map(|t| t.num_milliseconds() as u32); // XXX check for overflow

        //let mut statu

        {
            let statuses = try!(self.cp.get_many(&mut self.statuses[..], timeout));
            for status in statuses {
                let token = status.token();
                let overlapped = status.overlapped();
                let bytes_transferred = status.bytes_transferred();
                let handle = Handle { val: token };
                if self.observers[handle].read_overlapped == overlapped {
                    match self.observers[handle].read_fulfiller.take() {
                        None => (),
                        Some(f) => f.fulfill(bytes_transferred),
                    }
                }
                if self.observers[handle].write_overlapped == overlapped {
                    match self.observers[handle].write_fulfiller.take() {
                        None => (),
                        Some(f) => f.fulfill(bytes_transferred),
                    }
                }

            }
        }
        Ok(())
    }

    fn add_socket<T>(&mut self, sock: &T,
                     read_overlapped: *mut ::miow::Overlapped,
                     write_overlapped: *mut ::miow::Overlapped)
                     -> Result<Handle, ::std::io::Error>
        where T : ::std::os::windows::io::AsRawSocket
    {
        let observer = Observer {
            read_overlapped: read_overlapped,
            write_overlapped: write_overlapped,
            read_fulfiller: None,
            write_fulfiller: None,
        };
        let handle = self.observers.push(observer);
        try!(self.cp.add_socket(handle.val, sock));
        Ok(handle)
    }
}

#[derive(Clone)]
pub struct SocketAddressInner {
    reactor: Rc<RefCell<Reactor>>,
    addr: ::std::net::SocketAddr,
}

impl SocketAddressInner {
    pub fn new_tcp(reactor: Rc<RefCell<Reactor>>, addr: ::std::net::SocketAddr)
                   -> SocketAddressInner
    {
        SocketAddressInner {
            reactor: reactor,
            addr: addr,
        }
    }

    pub fn connect(&self) -> Promise<SocketStreamInner, ::std::io::Error> {
        use miow::net::TcpBuilderExt;
        let builder = match self.addr {
            ::std::net::SocketAddr::V4(_) => {
                let builder = pry!(::net2::TcpBuilder::new_v4());
                pry!(builder.bind("0.0.0.0:0"));
                builder
            }
            ::std::net::SocketAddr::V6(_) => {
                let builder = pry!(::net2::TcpBuilder::new_v6());
                pry!(builder.bind("[::]:0")); // TODO: does this work?
                builder
            }
        };

        let mut read_overlapped = Box::new(::miow::Overlapped::zero());
        let mut write_overlapped = Box::new(::miow::Overlapped::zero());
        let handle = pry!(
            self.reactor.borrow_mut().add_socket(&builder, &mut *read_overlapped,
                                                 &mut *write_overlapped));

        let (stream, _ready) = unsafe {
            pry!(builder.connect_overlapped(&self.addr, &mut *write_overlapped))
        };

        let reactor2 = self.reactor.clone();
        let result = self.reactor.borrow_mut().observers[handle].when_write_done().map(move |_| {
            Ok(SocketStreamInner::new_already_registered(
                reactor2, stream,
                read_overlapped, write_overlapped,
                handle))
        });

        result
    }

    pub fn listen(&mut self) -> Result<SocketListenerInner, ::std::io::Error> {
        let listener = try!(::std::net::TcpListener::bind(self.addr));
        SocketListenerInner::new(self.reactor.clone(), listener, self.addr)
    }
}

pub struct SocketListenerInner {
    reactor: Rc<RefCell<Reactor>>,
    listener: ::std::net::TcpListener,
    addr: ::std::net::SocketAddr,
    read_overlapped: Box<::miow::Overlapped>,
    handle: Handle,
    pub queue: Option<Promise<(),()>>,
}

impl SocketListenerInner {
    fn new(reactor: Rc<RefCell<Reactor>>, listener: ::std::net::TcpListener,
           addr: ::std::net::SocketAddr)
           -> Result<SocketListenerInner, ::std::io::Error>
    {
        let mut read_overlapped = Box::new(::miow::Overlapped::zero());
        let handle = try!(
            reactor.borrow_mut().add_socket(&listener, &mut *read_overlapped,
                                            ::std::ptr::null_mut()));

        Ok(SocketListenerInner {
            reactor: reactor,
            listener: listener,
            addr: addr,
            read_overlapped: read_overlapped,
            handle: handle,
            queue: None,
        })
    }

    pub fn accept_internal(inner: Rc<RefCell<SocketListenerInner>>)
                           -> Promise<SocketStreamInner, ::std::io::Error>
    {
        use miow::net::TcpListenerExt;
        let builder = match inner.borrow().addr {
            ::std::net::SocketAddr::V4(_) => pry!(::net2::TcpBuilder::new_v4()),
            ::std::net::SocketAddr::V6(_) => pry!(::net2::TcpBuilder::new_v6()),
        };

        let mut accept_addrs = ::miow::net::AcceptAddrsBuf::new();

        let &mut SocketListenerInner {
            ref reactor, ref mut listener, ref mut read_overlapped, handle,
            ..
        } =  &mut *inner.borrow_mut();

        let (stream, _ready) = unsafe {
            pry!(listener.accept_overlapped(&builder,
                                            &mut accept_addrs,
                                            read_overlapped))
        };

        let reactor2 = reactor.clone();
        let result = reactor.borrow_mut().observers[handle].when_read_done().map(move |_| {
            println!("accepted!");
            SocketStreamInner::new(reactor2, stream)
        });

        result
    }
}

pub struct SocketStreamInner {
    reactor: Rc<RefCell<Reactor>>,
    stream: ::std::net::TcpStream,
    read_overlapped: Box<::miow::Overlapped>,
    write_overlapped: Box<::miow::Overlapped>,
    handle: Handle,
    pub read_queue: Option<Promise<(),()>>,
    pub write_queue: Option<Promise<(),()>>,
}

impl SocketStreamInner {
    fn new(reactor: Rc<RefCell<Reactor>>, stream: ::std::net::TcpStream)
           -> Result<SocketStreamInner, ::std::io::Error>
    {
        let mut read_overlapped = Box::new(::miow::Overlapped::zero());
        let mut write_overlapped = Box::new(::miow::Overlapped::zero());
        let handle = try!(
            reactor.borrow_mut().add_socket(&stream, &mut *read_overlapped,
                                            &mut *write_overlapped));
        Ok(SocketStreamInner {
            reactor: reactor,
            stream: stream,
            read_overlapped: read_overlapped,
            write_overlapped: write_overlapped,
            handle: handle,
            read_queue: None,
            write_queue: None,
        })
    }

    fn new_already_registered(reactor: Rc<RefCell<Reactor>>, stream: ::std::net::TcpStream,
                              read_overlapped: Box<::miow::Overlapped>,
                              write_overlapped: Box<::miow::Overlapped>,
                              handle: Handle)
           -> SocketStreamInner
    {
        SocketStreamInner {
            reactor: reactor,
            stream: stream,
            read_overlapped: read_overlapped,
            write_overlapped: write_overlapped,
            handle: handle,
            read_queue: None,
            write_queue: None,
        }
    }

    pub fn try_read_internal<T>(inner: Rc<RefCell<SocketStreamInner>>,
                                mut buf: T,
                                already_read: usize,
                                min_bytes: usize)
                                -> Promise<(T, usize), ::std::io::Error>
        where T: AsMut<[u8]>
    {
        use ::miow::net::TcpStreamExt;

        if already_read >= min_bytes {
            return Promise::ok((buf, already_read));
        }

        let inner2 = inner.clone();
        let &mut SocketStreamInner {
            ref reactor, ref mut stream,
            ref mut read_overlapped, handle,
            ..
        } =  &mut *inner.borrow_mut();

        let _done = unsafe {
            pry!(stream.read_overlapped(&mut buf.as_mut()[already_read..], read_overlapped))
        };

        let result = reactor.borrow_mut().observers[handle].when_read_done().then(move |n| {
            println!("read transferred this many bytes: {}", n);
            let total_read = n as usize + already_read;
            if n == 0 {
                println!("EOF");
                Promise::ok((buf, total_read))
            } else {
                SocketStreamInner::try_read_internal(inner2, buf, total_read, min_bytes)
            }
        });

        result
    }

    pub fn write_internal<T>(inner: Rc<RefCell<SocketStreamInner>>,
                             buf: T,
                             already_written: usize) -> Promise<T, ::std::io::Error>
        where T: AsRef<[u8]>
    {
        use ::miow::net::TcpStreamExt;

        if already_written == buf.as_ref().len() {
            return Promise::ok(buf)
        }

        let inner2 = inner.clone();
        let &mut SocketStreamInner {
            ref reactor,
            ref stream,
            ref mut write_overlapped, handle,
            ..
        } =  &mut *inner.borrow_mut();

        let _done = unsafe {
            pry!(stream.write_overlapped(&buf.as_ref()[already_written ..], write_overlapped))
        };

        let result = reactor.borrow_mut().observers[handle].when_write_done().then(move |n| {
            println!("write transferred this many bytes: {}", n);
            let total_written = n as usize + already_written;
            if n == 0 {
                println!("wrote zero bytes!?");
            }

            SocketStreamInner::write_internal(inner2, buf, total_written)
        });

        result
    }
}
<|endoftext|>