<commit_before>// Copyright 2019 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.

//! A generic Serde-based `Transport` that can serialize anything supported by `tokio-serde` via any medium that implements `AsyncRead` and `AsyncWrite`.

#![deny(missing_docs)]

use futures::{prelude::*, task::*};
use pin_project::pin_project;
use serde::{Deserialize, Serialize};
use std::{error::Error, io, pin::Pin};
use tokio::io::{AsyncRead, AsyncWrite};
use tokio_serde::{Framed as SerdeFramed, *};
use tokio_util::codec::{
    length_delimited::{self, LengthDelimitedCodec},
    Framed,
};

/// A transport that serializes to, and deserializes from, a byte stream.
#[pin_project]
pub struct Transport<S, Item, SinkItem, Codec> {
    #[pin]
    inner: SerdeFramed<Framed<S, LengthDelimitedCodec>, Item, SinkItem, Codec>,
}

impl<S, Item, SinkItem, Codec> Transport<S, Item, SinkItem, Codec> {
    /// Returns the inner transport over which messages are sent and received.
    pub fn get_ref(&self) -> &S {
        self.inner.get_ref().get_ref()
    }
}

impl<S, Item, SinkItem, Codec, CodecError> Stream for Transport<S, Item, SinkItem, Codec>
where
    S: AsyncWrite + AsyncRead,
    Item: for<'a> Deserialize<'a>,
    Codec: Deserializer<Item>,
    CodecError: Into<Box<dyn std::error::Error + Send + Sync>>,
    SerdeFramed<Framed<S, LengthDelimitedCodec>, Item, SinkItem, Codec>:
        Stream<Item = Result<Item, CodecError>>,
{
    type Item = io::Result<Item>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<io::Result<Item>>> {
        match self.project().inner.poll_next(cx) {
            Poll::Pending => Poll::Pending,
            Poll::Ready(None) => Poll::Ready(None),
            Poll::Ready(Some(Ok::<_, CodecError>(next))) => Poll::Ready(Some(Ok(next))),
            Poll::Ready(Some(Err::<_, CodecError>(e))) => {
                Poll::Ready(Some(Err(io::Error::new(io::ErrorKind::Other, e))))
            }
        }
    }
}

impl<S, Item, SinkItem, Codec, CodecError> Sink<SinkItem> for Transport<S, Item, SinkItem, Codec>
where
    S: AsyncWrite,
    SinkItem: Serialize,
    Codec: Serializer<SinkItem>,
    CodecError: Into<Box<dyn Error + Send + Sync>>,
    SerdeFramed<Framed<S, LengthDelimitedCodec>, Item, SinkItem, Codec>:
        Sink<SinkItem, Error = CodecError>,
{
    type Error = io::Error;

    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        convert(self.project().inner.poll_ready(cx))
    }

    fn start_send(self: Pin<&mut Self>, item: SinkItem) -> io::Result<()> {
        self.project()
            .inner
            .start_send(item)
            .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        convert(self.project().inner.poll_flush(cx))
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        convert(self.project().inner.poll_close(cx))
    }
}

fn convert<E: Into<Box<dyn Error + Send + Sync>>>(
    poll: Poll<Result<(), E>>,
) -> Poll<io::Result<()>> {
    poll.map(|ready| ready.map_err(|e| io::Error::new(io::ErrorKind::Other, e)))
}

/// Constructs a new transport from a framed transport and a serialization codec.
pub fn new<S, Item, SinkItem, Codec>(
    framed_io: Framed<S, LengthDelimitedCodec>,
    codec: Codec,
) -> Transport<S, Item, SinkItem, Codec>
where
    S: AsyncWrite + AsyncRead,
    Item: for<'de> Deserialize<'de>,
    SinkItem: Serialize,
    Codec: Serializer<SinkItem> + Deserializer<Item>,
{
    Transport {
        inner: SerdeFramed::new(framed_io, codec),
    }
}

impl<S, Item, SinkItem, Codec> From<(S, Codec)> for Transport<S, Item, SinkItem, Codec>
where
    S: AsyncWrite + AsyncRead,
    Item: for<'de> Deserialize<'de>,
    SinkItem: Serialize,
    Codec: Serializer<SinkItem> + Deserializer<Item>,
{
    fn from((io, codec): (S, Codec)) -> Self {
        new(Framed::new(io, LengthDelimitedCodec::new()), codec)
    }
}

#[cfg(feature = "tcp")]
#[cfg_attr(docsrs, doc(cfg(feature = "tcp")))]
/// TCP support for generic transport using Tokio.
pub mod tcp {
    use {
        super::*,
        futures::ready,
        std::{marker::PhantomData, net::SocketAddr},
        tokio::net::{TcpListener, TcpStream, ToSocketAddrs},
    };

    mod private {
        use super::*;

        pub trait Sealed {}

        impl<Item, SinkItem, Codec> Sealed for Transport<TcpStream, Item, SinkItem, Codec> {}
    }

    impl<Item, SinkItem, Codec> Transport<TcpStream, Item, SinkItem, Codec> {
        /// Returns the peer address of the underlying TcpStream.
        pub fn peer_addr(&self) -> io::Result<SocketAddr> {
            self.inner.get_ref().get_ref().peer_addr()
        }
        /// Returns the local address of the underlying TcpStream.
        pub fn local_addr(&self) -> io::Result<SocketAddr> {
            self.inner.get_ref().get_ref().local_addr()
        }
    }

    /// A connection Future that also exposes the length-delimited framing config.
    #[pin_project]
    pub struct Connect<T, Item, SinkItem, CodecFn> {
        #[pin]
        inner: T,
        codec_fn: CodecFn,
        config: length_delimited::Builder,
        ghost: PhantomData<(fn(SinkItem), fn() -> Item)>,
    }

    impl<T, Item, SinkItem, Codec, CodecFn> Future for Connect<T, Item, SinkItem, CodecFn>
    where
        T: Future<Output = io::Result<TcpStream>>,
        Item: for<'de> Deserialize<'de>,
        SinkItem: Serialize,
        Codec: Serializer<SinkItem> + Deserializer<Item>,
        CodecFn: Fn() -> Codec,
    {
        type Output = io::Result<Transport<TcpStream, Item, SinkItem, Codec>>;

        fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
            let io = ready!(self.as_mut().project().inner.poll(cx))?;
            Poll::Ready(Ok(new(self.config.new_framed(io), (self.codec_fn)())))
        }
    }

    impl<T, Item, SinkItem, CodecFn> Connect<T, Item, SinkItem, CodecFn> {
        /// Returns an immutable reference to the length-delimited codec's config.
        pub fn config(&self) -> &length_delimited::Builder {
            &self.config
        }

        /// Returns a mutable reference to the length-delimited codec's config.
        pub fn config_mut(&mut self) -> &mut length_delimited::Builder {
            &mut self.config
        }
    }

    /// Connects to `addr`, wrapping the connection in a TCP transport.
    pub fn connect<A, Item, SinkItem, Codec, CodecFn>(
        addr: A,
        codec_fn: CodecFn,
    ) -> Connect<impl Future<Output = io::Result<TcpStream>>, Item, SinkItem, CodecFn>
    where
        A: ToSocketAddrs,
        Item: for<'de> Deserialize<'de>,
        SinkItem: Serialize,
        Codec: Serializer<SinkItem> + Deserializer<Item>,
        CodecFn: Fn() -> Codec,
    {
        Connect {
            inner: TcpStream::connect(addr),
            codec_fn,
            config: LengthDelimitedCodec::builder(),
            ghost: PhantomData,
        }
    }

    /// Listens on `addr`, wrapping accepted connections in TCP transports.
    pub async fn listen<A, Item, SinkItem, Codec, CodecFn>(
        addr: A,
        codec_fn: CodecFn,
    ) -> io::Result<Incoming<Item, SinkItem, Codec, CodecFn>>
    where
        A: ToSocketAddrs,
        Item: for<'de> Deserialize<'de>,
        Codec: Serializer<SinkItem> + Deserializer<Item>,
        CodecFn: Fn() -> Codec,
    {
        let listener = TcpListener::bind(addr).await?;
        let local_addr = listener.local_addr()?;
        Ok(Incoming {
            listener,
            codec_fn,
            local_addr,
            config: LengthDelimitedCodec::builder(),
            ghost: PhantomData,
        })
    }

    /// A [`TcpListener`] that wraps connections in [transports](Transport).
    #[pin_project]
    #[derive(Debug)]
    pub struct Incoming<Item, SinkItem, Codec, CodecFn> {
        listener: TcpListener,
        local_addr: SocketAddr,
        codec_fn: CodecFn,
        config: length_delimited::Builder,
        ghost: PhantomData<(fn() -> Item, fn(SinkItem), Codec)>,
    }

    impl<Item, SinkItem, Codec, CodecFn> Incoming<Item, SinkItem, Codec, CodecFn> {
        /// Returns the address being listened on.
        pub fn local_addr(&self) -> SocketAddr {
            self.local_addr
        }

        /// Returns an immutable reference to the length-delimited codec's config.
        pub fn config(&self) -> &length_delimited::Builder {
            &self.config
        }

        /// Returns a mutable reference to the length-delimited codec's config.
        pub fn config_mut(&mut self) -> &mut length_delimited::Builder {
            &mut self.config
        }
    }

    impl<Item, SinkItem, Codec, CodecFn> Stream for Incoming<Item, SinkItem, Codec, CodecFn>
    where
        Item: for<'de> Deserialize<'de>,
        SinkItem: Serialize,
        Codec: Serializer<SinkItem> + Deserializer<Item>,
        CodecFn: Fn() -> Codec,
    {
        type Item = io::Result<Transport<TcpStream, Item, SinkItem, Codec>>;

        fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
            let conn: TcpStream =
                ready!(Pin::new(&mut self.as_mut().project().listener).poll_accept(cx)?).0;
            Poll::Ready(Some(Ok(new(
                self.config.new_framed(conn),
                (self.codec_fn)(),
            ))))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::Transport;
    use assert_matches::assert_matches;
    use futures::{task::*, Sink, Stream};
    use pin_utils::pin_mut;
    use std::{
        io::{self, Cursor},
        pin::Pin,
    };
    use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
    use tokio_serde::formats::SymmetricalJson;

    fn ctx() -> Context<'static> {
        Context::from_waker(&noop_waker_ref())
    }

    #[test]
    fn test_stream() {
        struct TestIo(Cursor<&'static [u8]>);

        impl AsyncRead for TestIo {
            fn poll_read(
                mut self: Pin<&mut Self>,
                cx: &mut Context<'_>,
                buf: &mut ReadBuf<'_>,
            ) -> Poll<io::Result<()>> {
                AsyncRead::poll_read(Pin::new(self.0.get_mut()), cx, buf)
            }
        }

        impl AsyncWrite for TestIo {
            fn poll_write(
                self: Pin<&mut Self>,
                _cx: &mut Context<'_>,
                _buf: &[u8],
            ) -> Poll<io::Result<usize>> {
                unreachable!()
            }

            fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
                unreachable!()
            }

            fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
                unreachable!()
            }
        }

        let data = b"\x00\x00\x00\x18\"Test one, check check.\"";
        let transport = Transport::from((
            TestIo(Cursor::new(data)),
            SymmetricalJson::<String>::default(),
        ));
        pin_mut!(transport);

        assert_matches!(
            transport.poll_next(&mut ctx()),
            Poll::Ready(Some(Ok(ref s))) if s == "Test one, check check.");
    }

    #[test]
    fn test_sink() {
        struct TestIo<'a>(&'a mut Vec<u8>);

        impl<'a> AsyncRead for TestIo<'a> {
            fn poll_read(
                self: Pin<&mut Self>,
                _cx: &mut Context<'_>,
                _buf: &mut ReadBuf<'_>,
            ) -> Poll<io::Result<()>> {
                unreachable!()
            }
        }

        impl<'a> AsyncWrite for TestIo<'a> {
            fn poll_write(
                mut self: Pin<&mut Self>,
                cx: &mut Context<'_>,
                buf: &[u8],
            ) -> Poll<io::Result<usize>> {
                AsyncWrite::poll_write(Pin::new(&mut *self.0), cx, buf)
            }

            fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
                AsyncWrite::poll_flush(Pin::new(&mut *self.0), cx)
            }

            fn poll_shutdown(
                mut self: Pin<&mut Self>,
                cx: &mut Context<'_>,
            ) -> Poll<io::Result<()>> {
                AsyncWrite::poll_shutdown(Pin::new(&mut *self.0), cx)
            }
        }

        let mut writer = vec![];
        let transport =
            Transport::from((TestIo(&mut writer), SymmetricalJson::<String>::default()));
        pin_mut!(transport);

        assert_matches!(
            transport.as_mut().poll_ready(&mut ctx()),
            Poll::Ready(Ok(()))
        );
        assert_matches!(
            transport
                .as_mut()
                .start_send("Test one, check check.".into()),
            Ok(())
        );
        assert_matches!(transport.poll_flush(&mut ctx()), Poll::Ready(Ok(())));
        assert_eq!(writer, b"\x00\x00\x00\x18\"Test one, check check.\"");
    }
}
<commit_msg>Fix unused import<commit_after>// Copyright 2019 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.

//! A generic Serde-based `Transport` that can serialize anything supported by `tokio-serde` via any medium that implements `AsyncRead` and `AsyncWrite`.

#![deny(missing_docs)]

use futures::{prelude::*, task::*};
use pin_project::pin_project;
use serde::{Deserialize, Serialize};
use std::{error::Error, io, pin::Pin};
use tokio::io::{AsyncRead, AsyncWrite};
use tokio_serde::{Framed as SerdeFramed, *};
use tokio_util::codec::{
    length_delimited::LengthDelimitedCodec,
    Framed,
};

/// A transport that serializes to, and deserializes from, a byte stream.
#[pin_project]
pub struct Transport<S, Item, SinkItem, Codec> {
    #[pin]
    inner: SerdeFramed<Framed<S, LengthDelimitedCodec>, Item, SinkItem, Codec>,
}

impl<S, Item, SinkItem, Codec> Transport<S, Item, SinkItem, Codec> {
    /// Returns the inner transport over which messages are sent and received.
    pub fn get_ref(&self) -> &S {
        self.inner.get_ref().get_ref()
    }
}

impl<S, Item, SinkItem, Codec, CodecError> Stream for Transport<S, Item, SinkItem, Codec>
where
    S: AsyncWrite + AsyncRead,
    Item: for<'a> Deserialize<'a>,
    Codec: Deserializer<Item>,
    CodecError: Into<Box<dyn std::error::Error + Send + Sync>>,
    SerdeFramed<Framed<S, LengthDelimitedCodec>, Item, SinkItem, Codec>:
        Stream<Item = Result<Item, CodecError>>,
{
    type Item = io::Result<Item>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<io::Result<Item>>> {
        match self.project().inner.poll_next(cx) {
            Poll::Pending => Poll::Pending,
            Poll::Ready(None) => Poll::Ready(None),
            Poll::Ready(Some(Ok::<_, CodecError>(next))) => Poll::Ready(Some(Ok(next))),
            Poll::Ready(Some(Err::<_, CodecError>(e))) => {
                Poll::Ready(Some(Err(io::Error::new(io::ErrorKind::Other, e))))
            }
        }
    }
}

impl<S, Item, SinkItem, Codec, CodecError> Sink<SinkItem> for Transport<S, Item, SinkItem, Codec>
where
    S: AsyncWrite,
    SinkItem: Serialize,
    Codec: Serializer<SinkItem>,
    CodecError: Into<Box<dyn Error + Send + Sync>>,
    SerdeFramed<Framed<S, LengthDelimitedCodec>, Item, SinkItem, Codec>:
        Sink<SinkItem, Error = CodecError>,
{
    type Error = io::Error;

    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        convert(self.project().inner.poll_ready(cx))
    }

    fn start_send(self: Pin<&mut Self>, item: SinkItem) -> io::Result<()> {
        self.project()
            .inner
            .start_send(item)
            .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        convert(self.project().inner.poll_flush(cx))
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        convert(self.project().inner.poll_close(cx))
    }
}

fn convert<E: Into<Box<dyn Error + Send + Sync>>>(
    poll: Poll<Result<(), E>>,
) -> Poll<io::Result<()>> {
    poll.map(|ready| ready.map_err(|e| io::Error::new(io::ErrorKind::Other, e)))
}

/// Constructs a new transport from a framed transport and a serialization codec.
pub fn new<S, Item, SinkItem, Codec>(
    framed_io: Framed<S, LengthDelimitedCodec>,
    codec: Codec,
) -> Transport<S, Item, SinkItem, Codec>
where
    S: AsyncWrite + AsyncRead,
    Item: for<'de> Deserialize<'de>,
    SinkItem: Serialize,
    Codec: Serializer<SinkItem> + Deserializer<Item>,
{
    Transport {
        inner: SerdeFramed::new(framed_io, codec),
    }
}

impl<S, Item, SinkItem, Codec> From<(S, Codec)> for Transport<S, Item, SinkItem, Codec>
where
    S: AsyncWrite + AsyncRead,
    Item: for<'de> Deserialize<'de>,
    SinkItem: Serialize,
    Codec: Serializer<SinkItem> + Deserializer<Item>,
{
    fn from((io, codec): (S, Codec)) -> Self {
        new(Framed::new(io, LengthDelimitedCodec::new()), codec)
    }
}

#[cfg(feature = "tcp")]
#[cfg_attr(docsrs, doc(cfg(feature = "tcp")))]
/// TCP support for generic transport using Tokio.
pub mod tcp {
    use {
        super::*,
        futures::ready,
        std::{marker::PhantomData, net::SocketAddr},
        tokio::net::{TcpListener, TcpStream, ToSocketAddrs},
        tokio_util::codec::length_delimited,
    };

    mod private {
        use super::*;

        pub trait Sealed {}

        impl<Item, SinkItem, Codec> Sealed for Transport<TcpStream, Item, SinkItem, Codec> {}
    }

    impl<Item, SinkItem, Codec> Transport<TcpStream, Item, SinkItem, Codec> {
        /// Returns the peer address of the underlying TcpStream.
        pub fn peer_addr(&self) -> io::Result<SocketAddr> {
            self.inner.get_ref().get_ref().peer_addr()
        }
        /// Returns the local address of the underlying TcpStream.
        pub fn local_addr(&self) -> io::Result<SocketAddr> {
            self.inner.get_ref().get_ref().local_addr()
        }
    }

    /// A connection Future that also exposes the length-delimited framing config.
    #[pin_project]
    pub struct Connect<T, Item, SinkItem, CodecFn> {
        #[pin]
        inner: T,
        codec_fn: CodecFn,
        config: length_delimited::Builder,
        ghost: PhantomData<(fn(SinkItem), fn() -> Item)>,
    }

    impl<T, Item, SinkItem, Codec, CodecFn> Future for Connect<T, Item, SinkItem, CodecFn>
    where
        T: Future<Output = io::Result<TcpStream>>,
        Item: for<'de> Deserialize<'de>,
        SinkItem: Serialize,
        Codec: Serializer<SinkItem> + Deserializer<Item>,
        CodecFn: Fn() -> Codec,
    {
        type Output = io::Result<Transport<TcpStream, Item, SinkItem, Codec>>;

        fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
            let io = ready!(self.as_mut().project().inner.poll(cx))?;
            Poll::Ready(Ok(new(self.config.new_framed(io), (self.codec_fn)())))
        }
    }

    impl<T, Item, SinkItem, CodecFn> Connect<T, Item, SinkItem, CodecFn> {
        /// Returns an immutable reference to the length-delimited codec's config.
        pub fn config(&self) -> &length_delimited::Builder {
            &self.config
        }

        /// Returns a mutable reference to the length-delimited codec's config.
        pub fn config_mut(&mut self) -> &mut length_delimited::Builder {
            &mut self.config
        }
    }

    /// Connects to `addr`, wrapping the connection in a TCP transport.
    pub fn connect<A, Item, SinkItem, Codec, CodecFn>(
        addr: A,
        codec_fn: CodecFn,
    ) -> Connect<impl Future<Output = io::Result<TcpStream>>, Item, SinkItem, CodecFn>
    where
        A: ToSocketAddrs,
        Item: for<'de> Deserialize<'de>,
        SinkItem: Serialize,
        Codec: Serializer<SinkItem> + Deserializer<Item>,
        CodecFn: Fn() -> Codec,
    {
        Connect {
            inner: TcpStream::connect(addr),
            codec_fn,
            config: LengthDelimitedCodec::builder(),
            ghost: PhantomData,
        }
    }

    /// Listens on `addr`, wrapping accepted connections in TCP transports.
    pub async fn listen<A, Item, SinkItem, Codec, CodecFn>(
        addr: A,
        codec_fn: CodecFn,
    ) -> io::Result<Incoming<Item, SinkItem, Codec, CodecFn>>
    where
        A: ToSocketAddrs,
        Item: for<'de> Deserialize<'de>,
        Codec: Serializer<SinkItem> + Deserializer<Item>,
        CodecFn: Fn() -> Codec,
    {
        let listener = TcpListener::bind(addr).await?;
        let local_addr = listener.local_addr()?;
        Ok(Incoming {
            listener,
            codec_fn,
            local_addr,
            config: LengthDelimitedCodec::builder(),
            ghost: PhantomData,
        })
    }

    /// A [`TcpListener`] that wraps connections in [transports](Transport).
    #[pin_project]
    #[derive(Debug)]
    pub struct Incoming<Item, SinkItem, Codec, CodecFn> {
        listener: TcpListener,
        local_addr: SocketAddr,
        codec_fn: CodecFn,
        config: length_delimited::Builder,
        ghost: PhantomData<(fn() -> Item, fn(SinkItem), Codec)>,
    }

    impl<Item, SinkItem, Codec, CodecFn> Incoming<Item, SinkItem, Codec, CodecFn> {
        /// Returns the address being listened on.
        pub fn local_addr(&self) -> SocketAddr {
            self.local_addr
        }

        /// Returns an immutable reference to the length-delimited codec's config.
        pub fn config(&self) -> &length_delimited::Builder {
            &self.config
        }

        /// Returns a mutable reference to the length-delimited codec's config.
        pub fn config_mut(&mut self) -> &mut length_delimited::Builder {
            &mut self.config
        }
    }

    impl<Item, SinkItem, Codec, CodecFn> Stream for Incoming<Item, SinkItem, Codec, CodecFn>
    where
        Item: for<'de> Deserialize<'de>,
        SinkItem: Serialize,
        Codec: Serializer<SinkItem> + Deserializer<Item>,
        CodecFn: Fn() -> Codec,
    {
        type Item = io::Result<Transport<TcpStream, Item, SinkItem, Codec>>;

        fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
            let conn: TcpStream =
                ready!(Pin::new(&mut self.as_mut().project().listener).poll_accept(cx)?).0;
            Poll::Ready(Some(Ok(new(
                self.config.new_framed(conn),
                (self.codec_fn)(),
            ))))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::Transport;
    use assert_matches::assert_matches;
    use futures::{task::*, Sink, Stream};
    use pin_utils::pin_mut;
    use std::{
        io::{self, Cursor},
        pin::Pin,
    };
    use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
    use tokio_serde::formats::SymmetricalJson;

    fn ctx() -> Context<'static> {
        Context::from_waker(&noop_waker_ref())
    }

    #[test]
    fn test_stream() {
        struct TestIo(Cursor<&'static [u8]>);

        impl AsyncRead for TestIo {
            fn poll_read(
                mut self: Pin<&mut Self>,
                cx: &mut Context<'_>,
                buf: &mut ReadBuf<'_>,
            ) -> Poll<io::Result<()>> {
                AsyncRead::poll_read(Pin::new(self.0.get_mut()), cx, buf)
            }
        }

        impl AsyncWrite for TestIo {
            fn poll_write(
                self: Pin<&mut Self>,
                _cx: &mut Context<'_>,
                _buf: &[u8],
            ) -> Poll<io::Result<usize>> {
                unreachable!()
            }

            fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
                unreachable!()
            }

            fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
                unreachable!()
            }
        }

        let data = b"\x00\x00\x00\x18\"Test one, check check.\"";
        let transport = Transport::from((
            TestIo(Cursor::new(data)),
            SymmetricalJson::<String>::default(),
        ));
        pin_mut!(transport);

        assert_matches!(
            transport.poll_next(&mut ctx()),
            Poll::Ready(Some(Ok(ref s))) if s == "Test one, check check.");
    }

    #[test]
    fn test_sink() {
        struct TestIo<'a>(&'a mut Vec<u8>);

        impl<'a> AsyncRead for TestIo<'a> {
            fn poll_read(
                self: Pin<&mut Self>,
                _cx: &mut Context<'_>,
                _buf: &mut ReadBuf<'_>,
            ) -> Poll<io::Result<()>> {
                unreachable!()
            }
        }

        impl<'a> AsyncWrite for TestIo<'a> {
            fn poll_write(
                mut self: Pin<&mut Self>,
                cx: &mut Context<'_>,
                buf: &[u8],
            ) -> Poll<io::Result<usize>> {
                AsyncWrite::poll_write(Pin::new(&mut *self.0), cx, buf)
            }

            fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
                AsyncWrite::poll_flush(Pin::new(&mut *self.0), cx)
            }

            fn poll_shutdown(
                mut self: Pin<&mut Self>,
                cx: &mut Context<'_>,
            ) -> Poll<io::Result<()>> {
                AsyncWrite::poll_shutdown(Pin::new(&mut *self.0), cx)
            }
        }

        let mut writer = vec![];
        let transport =
            Transport::from((TestIo(&mut writer), SymmetricalJson::<String>::default()));
        pin_mut!(transport);

        assert_matches!(
            transport.as_mut().poll_ready(&mut ctx()),
            Poll::Ready(Ok(()))
        );
        assert_matches!(
            transport
                .as_mut()
                .start_send("Test one, check check.".into()),
            Ok(())
        );
        assert_matches!(transport.poll_flush(&mut ctx()), Poll::Ready(Ok(())));
        assert_eq!(writer, b"\x00\x00\x00\x18\"Test one, check check.\"");
    }
}
<|endoftext|>


<commit_before>//
// Sysinfo
//
// Copyright (c) 2018 Guillaume Gomez
//

use sys::component::{self, Component};
use sys::disk::Disk;
use sys::processor::*;
use sys::users::get_users;

use std::cell::UnsafeCell;
use std::collections::HashMap;
use std::mem::{size_of, zeroed, MaybeUninit};
use std::time::SystemTime;

use LoadAvg;
use Networks;
use Pid;
use ProcessExt;
use RefreshKind;
use SystemExt;
use User;

use windows::process::{
    compute_cpu_usage, get_handle, get_system_computation_time, update_disk_usage, update_memory,
    Process,
};
use windows::tools::*;

use ntapi::ntexapi::{
    NtQuerySystemInformation, SystemProcessInformation, SYSTEM_PROCESS_INFORMATION,
};
use winapi::shared::minwindef::{FALSE, TRUE};
use winapi::shared::ntdef::{PVOID, ULONG};
use winapi::shared::ntstatus::STATUS_INFO_LENGTH_MISMATCH;
use winapi::um::minwinbase::STILL_ACTIVE;
use winapi::um::processthreadsapi::GetExitCodeProcess;
use winapi::um::sysinfoapi::{GetTickCount64, GetVersionExW, GlobalMemoryStatusEx, MEMORYSTATUSEX};
use winapi::um::winnt::{HANDLE, OSVERSIONINFOW};

use utils::into_iter;

/// Struct containing the system's information.
pub struct System {
    process_list: HashMap<usize, Process>,
    mem_total: u64,
    mem_available: u64,
    swap_total: u64,
    swap_free: u64,
    global_processor: Processor,
    processors: Vec<Processor>,
    components: Vec<Component>,
    disks: Vec<Disk>,
    query: Option<Query>,
    networks: Networks,
    boot_time: u64,
    users: Vec<User>,
}

// Useful for parallel iterations.
struct Wrap<T>(T);

unsafe impl<T> Send for Wrap<T> {}
unsafe impl<T> Sync for Wrap<T> {}

unsafe fn boot_time() -> u64 {
    match SystemTime::now().duration_since(SystemTime::UNIX_EPOCH) {
        Ok(n) => n.as_secs() - GetTickCount64() / 1000,
        Err(_e) => {
            sysinfo_debug!("Failed to compute boot time: {:?}", _e);
            0
        }
    }
}

impl SystemExt for System {
    #[allow(non_snake_case)]
    fn new_with_specifics(refreshes: RefreshKind) -> System {
        let (processors, vendor_id, brand) = init_processors();
        let mut s = System {
            process_list: HashMap::with_capacity(500),
            mem_total: 0,
            mem_available: 0,
            swap_total: 0,
            swap_free: 0,
            global_processor: Processor::new_with_values("Total CPU", vendor_id, brand, 0),
            processors,
            components: Vec::new(),
            disks: Vec::with_capacity(2),
            query: Query::new(),
            networks: Networks::new(),
            boot_time: unsafe { boot_time() },
            users: Vec::new(),
        };
        // TODO: in case a translation fails, it might be nice to log it somewhere...
        if let Some(ref mut query) = s.query {
            let x = unsafe { load_symbols() };
            if let Some(processor_trans) = get_translation(&"Processor".to_owned(), &x) {
                // let idle_time_trans = get_translation(&"% Idle Time".to_owned(), &x);
                let proc_time_trans = get_translation(&"% Processor Time".to_owned(), &x);
                if let Some(ref proc_time_trans) = proc_time_trans {
                    add_counter(
                        format!("\\{}(_Total)\\{}", processor_trans, proc_time_trans),
                        query,
                        get_key_used(&mut s.global_processor),
                        "tot_0".to_owned(),
                    );
                }
                for (pos, proc_) in s.processors.iter_mut().enumerate() {
                    if let Some(ref proc_time_trans) = proc_time_trans {
                        add_counter(
                            format!("\\{}({})\\{}", processor_trans, pos, proc_time_trans),
                            query,
                            get_key_used(proc_),
                            format!("{}_0", pos),
                        );
                    }
                }
            } else {
                sysinfo_debug!("failed to get `Processor` translation");
            }
        }
        s.refresh_specifics(refreshes);
        s
    }

    fn refresh_cpu(&mut self) {
        if let Some(ref mut query) = self.query {
            query.refresh();
            let mut used_time = None;
            if let Some(ref key_used) = *get_key_used(&mut self.global_processor) {
                used_time = Some(
                    query
                        .get(&key_used.unique_id)
                        .expect("global_key_idle disappeared"),
                );
            }
            if let Some(used_time) = used_time {
                self.global_processor.set_cpu_usage(used_time);
            }
            for p in self.processors.iter_mut() {
                let mut used_time = None;
                if let Some(ref key_used) = *get_key_used(p) {
                    used_time = Some(
                        query
                            .get(&key_used.unique_id)
                            .expect("key_used disappeared"),
                    );
                }
                if let Some(used_time) = used_time {
                    p.set_cpu_usage(used_time);
                }
            }
        }
    }

    fn refresh_memory(&mut self) {
        unsafe {
            let mut mem_info: MEMORYSTATUSEX = zeroed();
            mem_info.dwLength = size_of::<MEMORYSTATUSEX>() as u32;
            GlobalMemoryStatusEx(&mut mem_info);
            self.mem_total = auto_cast!(mem_info.ullTotalPhys, u64) / 1_000;
            self.mem_available = auto_cast!(mem_info.ullAvailPhys, u64) / 1_000;
            //self.swap_total = auto_cast!(mem_info.ullTotalPageFile - mem_info.ullTotalPhys, u64);
            //self.swap_free = auto_cast!(mem_info.ullAvailPageFile, u64);
        }
    }

    fn refresh_components_list(&mut self) {
        self.components = component::get_components();
    }

    #[allow(clippy::map_entry)]
    fn refresh_process(&mut self, pid: Pid) -> bool {
        if self.process_list.contains_key(&pid) {
            if !refresh_existing_process(self, pid) {
                self.process_list.remove(&pid);
                return false;
            }
            true
        } else if let Some(mut p) = Process::new_from_pid(pid) {
            let system_time = get_system_computation_time();
            compute_cpu_usage(&mut p, self.processors.len() as u64, system_time);
            update_disk_usage(&mut p);
            self.process_list.insert(pid, p);
            true
        } else {
            false
        }
    }

    #[allow(clippy::cast_ptr_alignment)]
    fn refresh_processes(&mut self) {
        // Windows 10 notebook requires at least 512KiB of memory to make it in one go
        let mut buffer_size: usize = 512 * 1024;

        loop {
            let mut process_information: Vec<u8> = Vec::with_capacity(buffer_size);

            let mut cb_needed = 0;
            let ntstatus = unsafe {
                process_information.set_len(buffer_size);
                NtQuerySystemInformation(
                    SystemProcessInformation,
                    process_information.as_mut_ptr() as PVOID,
                    buffer_size as ULONG,
                    &mut cb_needed,
                )
            };

            if ntstatus != STATUS_INFO_LENGTH_MISMATCH {
                if ntstatus < 0 {
                    sysinfo_debug!(
                        "Couldn't get process infos: NtQuerySystemInformation returned {}",
                        ntstatus
                    );
                }

                // Parse the data block to get process information
                let mut process_ids = Vec::with_capacity(500);
                let mut process_information_offset = 0;
                loop {
                    let p = unsafe {
                        process_information
                            .as_ptr()
                            .offset(process_information_offset)
                            as *const SYSTEM_PROCESS_INFORMATION
                    };
                    let pi = unsafe { &*p };

                    process_ids.push(Wrap(p));

                    if pi.NextEntryOffset == 0 {
                        break;
                    }

                    process_information_offset += pi.NextEntryOffset as isize;
                }
                let nb_processors = self.processors.len() as u64;
                let process_list = Wrap(UnsafeCell::new(&mut self.process_list));
                let system_time = get_system_computation_time();

                #[cfg(feature = "multithread")]
                use rayon::iter::ParallelIterator;

                // TODO: instead of using parallel iterator only here, would be better to be able
                //       to run it over `process_information` directly!
                let processes = into_iter(process_ids)
                    .filter_map(|pi| unsafe {
                        let pi = *pi.0;
                        let pid = pi.UniqueProcessId as usize;
                        if let Some(proc_) = (*process_list.0.get()).get_mut(&pid) {
                            proc_.memory = (pi.WorkingSetSize as u64) / 1_000;
                            proc_.virtual_memory = (pi.VirtualSize as u64) / 1_000;
                            compute_cpu_usage(proc_, nb_processors, system_time);
                            update_disk_usage(proc_);
                            proc_.updated = true;
                            return None;
                        }
                        let name = get_process_name(&pi, pid);
                        let mut p = Process::new_full(
                            pid,
                            if pi.InheritedFromUniqueProcessId as usize != 0 {
                                Some(pi.InheritedFromUniqueProcessId as usize)
                            } else {
                                None
                            },
                            (pi.WorkingSetSize as u64) / 1_000,
                            (pi.VirtualSize as u64) / 1_000,
                            name,
                        );
                        compute_cpu_usage(&mut p, nb_processors, system_time);
                        update_disk_usage(&mut p);
                        Some(p)
                    })
                    .collect::<Vec<_>>();
                self.process_list.retain(|_, v| {
                    let x = v.updated;
                    v.updated = false;
                    x
                });
                for p in processes.into_iter() {
                    self.process_list.insert(p.pid(), p);
                }

                break;
            }

            // GetNewBufferSize
            if cb_needed == 0 {
                buffer_size *= 2;
                continue;
            }
            // allocating a few more kilo bytes just in case there are some new process
            // kicked in since new call to NtQuerySystemInformation
            buffer_size = (cb_needed + (1024 * 10)) as usize;
        }
    }

    fn refresh_disks_list(&mut self) {
        self.disks = unsafe { get_disks() };
    }

    fn refresh_users_list(&mut self) {
        self.users = unsafe { get_users() };
    }

    fn get_processes(&self) -> &HashMap<Pid, Process> {
        &self.process_list
    }

    fn get_process(&self, pid: Pid) -> Option<&Process> {
        self.process_list.get(&(pid as usize))
    }

    fn get_global_processor_info(&self) -> &Processor {
        &self.global_processor
    }

    fn get_processors(&self) -> &[Processor] {
        &self.processors
    }

    fn get_total_memory(&self) -> u64 {
        self.mem_total
    }

    fn get_free_memory(&self) -> u64 {
        // MEMORYSTATUSEX doesn't report free memory
        self.mem_available
    }

    fn get_available_memory(&self) -> u64 {
        self.mem_available
    }

    fn get_used_memory(&self) -> u64 {
        self.mem_total - self.mem_available
    }

    fn get_total_swap(&self) -> u64 {
        self.swap_total
    }

    fn get_free_swap(&self) -> u64 {
        self.swap_free
    }

    fn get_used_swap(&self) -> u64 {
        self.swap_total - self.swap_free
    }

    fn get_components(&self) -> &[Component] {
        &self.components
    }

    fn get_components_mut(&mut self) -> &mut [Component] {
        &mut self.components
    }

    fn get_disks(&self) -> &[Disk] {
        &self.disks
    }

    fn get_disks_mut(&mut self) -> &mut [Disk] {
        &mut self.disks
    }

    fn get_users(&self) -> &[User] {
        &self.users
    }

    fn get_networks(&self) -> &Networks {
        &self.networks
    }

    fn get_networks_mut(&mut self) -> &mut Networks {
        &mut self.networks
    }

    fn get_uptime(&self) -> u64 {
        unsafe { GetTickCount64() / 1000 }
    }

    fn get_boot_time(&self) -> u64 {
        self.boot_time
    }

    fn get_load_average(&self) -> LoadAvg {
        get_load_average()
    }

    fn get_name(&self) -> Option<String> {
        Some("Windows".to_owned())
    }

    fn get_version(&self) -> Option<String> {
        get_system_version()
    }
}

impl Default for System {
    fn default() -> System {
        System::new()
    }
}

fn is_proc_running(handle: HANDLE) -> bool {
    let mut exit_code = 0;
    let ret = unsafe { GetExitCodeProcess(handle, &mut exit_code) };
    !(ret == FALSE || exit_code != STILL_ACTIVE)
}

fn refresh_existing_process(s: &mut System, pid: Pid) -> bool {
    if let Some(ref mut entry) = s.process_list.get_mut(&(pid as usize)) {
        if !is_proc_running(get_handle(entry)) {
            return false;
        }
        update_memory(entry);
        update_disk_usage(entry);
        compute_cpu_usage(
            entry,
            s.processors.len() as u64,
            get_system_computation_time(),
        );
        true
    } else {
        false
    }
}

#[allow(clippy::size_of_in_element_count)]
//^ needed for "name.Length as usize / std::mem::size_of::<u16>()"
pub(crate) fn get_process_name(process: &SYSTEM_PROCESS_INFORMATION, process_id: usize) -> String {
    let name = &process.ImageName;
    if name.Buffer.is_null() {
        match process_id {
            0 => "Idle".to_owned(),
            4 => "System".to_owned(),
            _ => format!("<no name> Process {}", process_id),
        }
    } else {
        let slice = unsafe {
            std::slice::from_raw_parts(
                name.Buffer,
                // The length is in bytes, not the length of string
                name.Length as usize / std::mem::size_of::<u16>(),
            )
        };

        String::from_utf16_lossy(slice)
    }
}

fn get_system_version() -> Option<String> {
    let mut info: OSVERSIONINFOW = unsafe { MaybeUninit::zeroed().assume_init() };
    info.dwOSVersionInfoSize = size_of::<OSVERSIONINFOW>() as _;
    if unsafe { GetVersionExW(&mut info) } == TRUE {
        Some(format!(
            "{}.{}.{}",
            info.dwMajorVersion, info.dwMinorVersion, info.dwBuildNumber
        ))
    } else {
        None
    }
}
<commit_msg>Working windows implementation<commit_after>//
// Sysinfo
//
// Copyright (c) 2018 Guillaume Gomez
//

use sys::component::{self, Component};
use sys::disk::Disk;
use sys::processor::*;
use sys::users::get_users;

use std::cell::UnsafeCell;
use std::collections::HashMap;
use std::mem::{size_of, zeroed, MaybeUninit};
use std::time::SystemTime;

use LoadAvg;
use Networks;
use Pid;
use ProcessExt;
use RefreshKind;
use SystemExt;
use User;

use windows::process::{
    compute_cpu_usage, get_handle, get_system_computation_time, update_disk_usage, update_memory,
    Process,
};
use windows::tools::*;

use ntapi::ntexapi::{
    NtQuerySystemInformation, SystemProcessInformation, SYSTEM_PROCESS_INFORMATION,
};
use winapi::ctypes::wchar_t;
use winapi::shared::minwindef::{FALSE, TRUE};
use winapi::shared::ntdef::{PVOID, ULONG};
use winapi::shared::ntstatus::STATUS_INFO_LENGTH_MISMATCH;
use winapi::um::minwinbase::STILL_ACTIVE;
use winapi::um::processthreadsapi::GetExitCodeProcess;
use winapi::um::sysinfoapi::{ComputerNamePhysicalDnsHostname, GetComputerNameExW, GetTickCount64, GetVersionExW, GlobalMemoryStatusEx, MEMORYSTATUSEX};
use winapi::um::winnt::{HANDLE, OSVERSIONINFOW};

use utils::into_iter;

/// Struct containing the system's information.
pub struct System {
    process_list: HashMap<usize, Process>,
    mem_total: u64,
    mem_available: u64,
    swap_total: u64,
    swap_free: u64,
    global_processor: Processor,
    processors: Vec<Processor>,
    components: Vec<Component>,
    disks: Vec<Disk>,
    query: Option<Query>,
    networks: Networks,
    boot_time: u64,
    users: Vec<User>,
}

// Useful for parallel iterations.
struct Wrap<T>(T);

unsafe impl<T> Send for Wrap<T> {}
unsafe impl<T> Sync for Wrap<T> {}

unsafe fn boot_time() -> u64 {
    match SystemTime::now().duration_since(SystemTime::UNIX_EPOCH) {
        Ok(n) => n.as_secs() - GetTickCount64() / 1000,
        Err(_e) => {
            sysinfo_debug!("Failed to compute boot time: {:?}", _e);
            0
        }
    }
}

impl SystemExt for System {
    #[allow(non_snake_case)]
    fn new_with_specifics(refreshes: RefreshKind) -> System {
        let (processors, vendor_id, brand) = init_processors();
        let mut s = System {
            process_list: HashMap::with_capacity(500),
            mem_total: 0,
            mem_available: 0,
            swap_total: 0,
            swap_free: 0,
            global_processor: Processor::new_with_values("Total CPU", vendor_id, brand, 0),
            processors,
            components: Vec::new(),
            disks: Vec::with_capacity(2),
            query: Query::new(),
            networks: Networks::new(),
            boot_time: unsafe { boot_time() },
            users: Vec::new(),
        };
        // TODO: in case a translation fails, it might be nice to log it somewhere...
        if let Some(ref mut query) = s.query {
            let x = unsafe { load_symbols() };
            if let Some(processor_trans) = get_translation(&"Processor".to_owned(), &x) {
                // let idle_time_trans = get_translation(&"% Idle Time".to_owned(), &x);
                let proc_time_trans = get_translation(&"% Processor Time".to_owned(), &x);
                if let Some(ref proc_time_trans) = proc_time_trans {
                    add_counter(
                        format!("\\{}(_Total)\\{}", processor_trans, proc_time_trans),
                        query,
                        get_key_used(&mut s.global_processor),
                        "tot_0".to_owned(),
                    );
                }
                for (pos, proc_) in s.processors.iter_mut().enumerate() {
                    if let Some(ref proc_time_trans) = proc_time_trans {
                        add_counter(
                            format!("\\{}({})\\{}", processor_trans, pos, proc_time_trans),
                            query,
                            get_key_used(proc_),
                            format!("{}_0", pos),
                        );
                    }
                }
            } else {
                sysinfo_debug!("failed to get `Processor` translation");
            }
        }
        s.refresh_specifics(refreshes);
        s
    }

    fn refresh_cpu(&mut self) {
        if let Some(ref mut query) = self.query {
            query.refresh();
            let mut used_time = None;
            if let Some(ref key_used) = *get_key_used(&mut self.global_processor) {
                used_time = Some(
                    query
                        .get(&key_used.unique_id)
                        .expect("global_key_idle disappeared"),
                );
            }
            if let Some(used_time) = used_time {
                self.global_processor.set_cpu_usage(used_time);
            }
            for p in self.processors.iter_mut() {
                let mut used_time = None;
                if let Some(ref key_used) = *get_key_used(p) {
                    used_time = Some(
                        query
                            .get(&key_used.unique_id)
                            .expect("key_used disappeared"),
                    );
                }
                if let Some(used_time) = used_time {
                    p.set_cpu_usage(used_time);
                }
            }
        }
    }

    fn refresh_memory(&mut self) {
        unsafe {
            let mut mem_info: MEMORYSTATUSEX = zeroed();
            mem_info.dwLength = size_of::<MEMORYSTATUSEX>() as u32;
            GlobalMemoryStatusEx(&mut mem_info);
            self.mem_total = auto_cast!(mem_info.ullTotalPhys, u64) / 1_000;
            self.mem_available = auto_cast!(mem_info.ullAvailPhys, u64) / 1_000;
            //self.swap_total = auto_cast!(mem_info.ullTotalPageFile - mem_info.ullTotalPhys, u64);
            //self.swap_free = auto_cast!(mem_info.ullAvailPageFile, u64);
        }
    }

    fn refresh_components_list(&mut self) {
        self.components = component::get_components();
    }

    #[allow(clippy::map_entry)]
    fn refresh_process(&mut self, pid: Pid) -> bool {
        if self.process_list.contains_key(&pid) {
            if !refresh_existing_process(self, pid) {
                self.process_list.remove(&pid);
                return false;
            }
            true
        } else if let Some(mut p) = Process::new_from_pid(pid) {
            let system_time = get_system_computation_time();
            compute_cpu_usage(&mut p, self.processors.len() as u64, system_time);
            update_disk_usage(&mut p);
            self.process_list.insert(pid, p);
            true
        } else {
            false
        }
    }

    #[allow(clippy::cast_ptr_alignment)]
    fn refresh_processes(&mut self) {
        // Windows 10 notebook requires at least 512KiB of memory to make it in one go
        let mut buffer_size: usize = 512 * 1024;

        loop {
            let mut process_information: Vec<u8> = Vec::with_capacity(buffer_size);

            let mut cb_needed = 0;
            let ntstatus = unsafe {
                process_information.set_len(buffer_size);
                NtQuerySystemInformation(
                    SystemProcessInformation,
                    process_information.as_mut_ptr() as PVOID,
                    buffer_size as ULONG,
                    &mut cb_needed,
                )
            };

            if ntstatus != STATUS_INFO_LENGTH_MISMATCH {
                if ntstatus < 0 {
                    sysinfo_debug!(
                        "Couldn't get process infos: NtQuerySystemInformation returned {}",
                        ntstatus
                    );
                }

                // Parse the data block to get process information
                let mut process_ids = Vec::with_capacity(500);
                let mut process_information_offset = 0;
                loop {
                    let p = unsafe {
                        process_information
                            .as_ptr()
                            .offset(process_information_offset)
                            as *const SYSTEM_PROCESS_INFORMATION
                    };
                    let pi = unsafe { &*p };

                    process_ids.push(Wrap(p));

                    if pi.NextEntryOffset == 0 {
                        break;
                    }

                    process_information_offset += pi.NextEntryOffset as isize;
                }
                let nb_processors = self.processors.len() as u64;
                let process_list = Wrap(UnsafeCell::new(&mut self.process_list));
                let system_time = get_system_computation_time();

                #[cfg(feature = "multithread")]
                use rayon::iter::ParallelIterator;

                // TODO: instead of using parallel iterator only here, would be better to be able
                //       to run it over `process_information` directly!
                let processes = into_iter(process_ids)
                    .filter_map(|pi| unsafe {
                        let pi = *pi.0;
                        let pid = pi.UniqueProcessId as usize;
                        if let Some(proc_) = (*process_list.0.get()).get_mut(&pid) {
                            proc_.memory = (pi.WorkingSetSize as u64) / 1_000;
                            proc_.virtual_memory = (pi.VirtualSize as u64) / 1_000;
                            compute_cpu_usage(proc_, nb_processors, system_time);
                            update_disk_usage(proc_);
                            proc_.updated = true;
                            return None;
                        }
                        let name = get_process_name(&pi, pid);
                        let mut p = Process::new_full(
                            pid,
                            if pi.InheritedFromUniqueProcessId as usize != 0 {
                                Some(pi.InheritedFromUniqueProcessId as usize)
                            } else {
                                None
                            },
                            (pi.WorkingSetSize as u64) / 1_000,
                            (pi.VirtualSize as u64) / 1_000,
                            name,
                        );
                        compute_cpu_usage(&mut p, nb_processors, system_time);
                        update_disk_usage(&mut p);
                        Some(p)
                    })
                    .collect::<Vec<_>>();
                self.process_list.retain(|_, v| {
                    let x = v.updated;
                    v.updated = false;
                    x
                });
                for p in processes.into_iter() {
                    self.process_list.insert(p.pid(), p);
                }

                break;
            }

            // GetNewBufferSize
            if cb_needed == 0 {
                buffer_size *= 2;
                continue;
            }
            // allocating a few more kilo bytes just in case there are some new process
            // kicked in since new call to NtQuerySystemInformation
            buffer_size = (cb_needed + (1024 * 10)) as usize;
        }
    }

    fn refresh_disks_list(&mut self) {
        self.disks = unsafe { get_disks() };
    }

    fn refresh_users_list(&mut self) {
        self.users = unsafe { get_users() };
    }

    fn get_processes(&self) -> &HashMap<Pid, Process> {
        &self.process_list
    }

    fn get_process(&self, pid: Pid) -> Option<&Process> {
        self.process_list.get(&(pid as usize))
    }

    fn get_global_processor_info(&self) -> &Processor {
        &self.global_processor
    }

    fn get_processors(&self) -> &[Processor] {
        &self.processors
    }

    fn get_total_memory(&self) -> u64 {
        self.mem_total
    }

    fn get_free_memory(&self) -> u64 {
        // MEMORYSTATUSEX doesn't report free memory
        self.mem_available
    }

    fn get_available_memory(&self) -> u64 {
        self.mem_available
    }

    fn get_used_memory(&self) -> u64 {
        self.mem_total - self.mem_available
    }

    fn get_total_swap(&self) -> u64 {
        self.swap_total
    }

    fn get_free_swap(&self) -> u64 {
        self.swap_free
    }

    fn get_used_swap(&self) -> u64 {
        self.swap_total - self.swap_free
    }

    fn get_components(&self) -> &[Component] {
        &self.components
    }

    fn get_components_mut(&mut self) -> &mut [Component] {
        &mut self.components
    }

    fn get_disks(&self) -> &[Disk] {
        &self.disks
    }

    fn get_disks_mut(&mut self) -> &mut [Disk] {
        &mut self.disks
    }

    fn get_users(&self) -> &[User] {
        &self.users
    }

    fn get_networks(&self) -> &Networks {
        &self.networks
    }

    fn get_networks_mut(&mut self) -> &mut Networks {
        &mut self.networks
    }

    fn get_uptime(&self) -> u64 {
        unsafe { GetTickCount64() / 1000 }
    }

    fn get_boot_time(&self) -> u64 {
        self.boot_time
    }

    fn get_load_average(&self) -> LoadAvg {
        get_load_average()
    }

    fn get_name(&self) -> Option<String> {
        Some("Windows".to_owned())
    }

    fn get_host_name(&self) -> Option<String> {
        get_dns_hostname()
    }

    fn get_version(&self) -> Option<String> {
        get_system_version()
    }
}

impl Default for System {
    fn default() -> System {
        System::new()
    }
}

fn is_proc_running(handle: HANDLE) -> bool {
    let mut exit_code = 0;
    let ret = unsafe { GetExitCodeProcess(handle, &mut exit_code) };
    !(ret == FALSE || exit_code != STILL_ACTIVE)
}

fn refresh_existing_process(s: &mut System, pid: Pid) -> bool {
    if let Some(ref mut entry) = s.process_list.get_mut(&(pid as usize)) {
        if !is_proc_running(get_handle(entry)) {
            return false;
        }
        update_memory(entry);
        update_disk_usage(entry);
        compute_cpu_usage(
            entry,
            s.processors.len() as u64,
            get_system_computation_time(),
        );
        true
    } else {
        false
    }
}

#[allow(clippy::size_of_in_element_count)]
//^ needed for "name.Length as usize / std::mem::size_of::<u16>()"
pub(crate) fn get_process_name(process: &SYSTEM_PROCESS_INFORMATION, process_id: usize) -> String {
    let name = &process.ImageName;
    if name.Buffer.is_null() {
        match process_id {
            0 => "Idle".to_owned(),
            4 => "System".to_owned(),
            _ => format!("<no name> Process {}", process_id),
        }
    } else {
        let slice = unsafe {
            std::slice::from_raw_parts(
                name.Buffer,
                // The length is in bytes, not the length of string
                name.Length as usize / std::mem::size_of::<u16>(),
            )
        };

        String::from_utf16_lossy(slice)
    }
}

fn get_system_version() -> Option<String> {
    let mut info: OSVERSIONINFOW = unsafe { MaybeUninit::zeroed().assume_init() };
    info.dwOSVersionInfoSize = size_of::<OSVERSIONINFOW>() as _;
    if unsafe { GetVersionExW(&mut info) } == TRUE {
        Some(format!(
            "{}.{}.{}",
            info.dwMajorVersion, info.dwMinorVersion, info.dwBuildNumber
        ))
    } else {
        None
    }
}

fn get_dns_hostname() -> Option<String> {
    let mut buffer_size = 0;
    // Running this first to get the buffer size since the DNS name can be longer than MAX_COMPUTERNAME_LENGTH
    // setting the `lpBuffer` to null will return the buffer size
    // https://docs.microsoft.com/en-us/windows/win32/api/sysinfoapi/nf-sysinfoapi-getcomputernameexw
    unsafe{GetComputerNameExW(ComputerNamePhysicalDnsHostname, std::ptr::null_mut(), &mut buffer_size)};

    // Setting the buffer with the new length
    let mut buffer = vec![0 as wchar_t; buffer_size as usize];

    // https://docs.microsoft.com/en-us/windows/win32/api/sysinfoapi/ne-sysinfoapi-computer_name_format
    if unsafe {GetComputerNameExW(ComputerNamePhysicalDnsHostname, buffer.as_mut_ptr() as *mut wchar_t, &mut buffer_size)} == TRUE {
        String::from_utf16(&buffer).ok()
    } else {
        sysinfo_debug!("Failed to get computer hostname");
        None
    }
}<|endoftext|>


<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use fnv::{FnvHashMap, FnvHashSet};

use graphql_ir::*;
use interner::StringKey;
use std::collections::hash_map::Entry;
use std::fmt;

struct Node {
    ir: Option<ExecutableDefinition>,
    parents: Vec<StringKey>,
    children: Vec<StringKey>,
}

impl fmt::Debug for Node {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_fmt(format_args!(
            "parents: {:?} / children: {:?}",
            self.parents, self.children
        ))
    }
}

pub fn get_reachable_ir(
    definitions: Vec<ExecutableDefinition>,
    base_definition_names: FnvHashSet<StringKey>,
    reachable_names: FnvHashSet<StringKey>,
) -> Vec<ExecutableDefinition> {
    if reachable_names.is_empty() {
        return vec![];
    }

    let trees = build_dependency_trees(definitions);

    let mut visited = FnvHashSet::default();
    let mut filtered_definitions = FnvHashMap::default();

    for key in reachable_names.into_iter() {
        if trees.contains_key(&key) {
            add_related_nodes(
                &mut visited,
                &mut filtered_definitions,
                &trees,
                &base_definition_names,
                key,
            );
        }
    }

    filtered_definitions
        .drain()
        .map(|(_, definition)| definition)
        .collect()
}

// Build a set of dependency trees that nodes are "doubly linked"
fn build_dependency_trees(definitions: Vec<ExecutableDefinition>) -> FnvHashMap<StringKey, Node> {
    let mut trees = FnvHashMap::with_capacity_and_hasher(definitions.len(), Default::default());

    for definition in definitions.into_iter() {
        let name = match &definition {
            ExecutableDefinition::Operation(operation) => operation.name.item,
            ExecutableDefinition::Fragment(fragment) => fragment.name.item,
        };

        // Visit the selections of the IR to build it's `children`
        let mut children = vec![];
        let selections = match &definition {
            ExecutableDefinition::Operation(operation) => &operation.selections,
            ExecutableDefinition::Fragment(fragment) => &fragment.selections,
        };
        visit_selections(&mut trees, &selections, name, &mut children);

        // Insert or update the representation of the IR in the dependency tree
        match trees.entry(name) {
            // Add a new node for current IR to the dependency tree
            Entry::Vacant(entry) => {
                entry.insert(Node {
                    ir: Some(definition),
                    parents: vec![],
                    children,
                });
            }

            // The node is already created when visiting selections of the IR, but it's `ir` and `children` haven't been set
            Entry::Occupied(mut entry) => {
                let node = entry.get_mut();
                if let Some(def) = &node.ir {
                    panic!(
                        "Duplicate definition: had {:?} and found another {:?}.",
                        def, node.ir
                    );
                }
                node.ir = Some(definition);
                node.children = children;
            }
        }
    }
    trees
}

// Visit the selections of current IR, set the `children` for the node representing the IR,
// and the `parents` for nodes representing the children IR
fn visit_selections(
    trees: &mut FnvHashMap<StringKey, Node>,
    selections: &[Selection],
    parent_name: StringKey,
    children: &mut Vec<StringKey>,
) {
    for selection in selections {
        match selection {
            Selection::FragmentSpread(node) => {
                let key = node.fragment.item;
                match trees.get_mut(&key) {
                    None => {
                        trees.insert(
                            key,
                            Node {
                                ir: None,
                                parents: vec![parent_name],
                                children: vec![],
                            },
                        );
                    }
                    Some(node) => {
                        node.parents.push(parent_name);
                    }
                }
                children.push(key);
            }
            Selection::LinkedField(node) => {
                visit_selections(trees, &node.selections, parent_name, children);
            }
            Selection::InlineFragment(node) => {
                visit_selections(trees, &node.selections, parent_name, children);
            }
            Selection::ScalarField(_) => {}
            Selection::Condition(node) => {
                visit_selections(trees, &node.selections, parent_name, children);
            }
        }
    }
}

// From `key` of changed definition, recusively traverse up the depenency tree, and add all related nodes (ancestors
// of changned definitions which are not from base definitions, and all of their desendants) into the `result`
fn add_related_nodes(
    visited: &mut FnvHashSet<StringKey>,
    result: &mut FnvHashMap<StringKey, ExecutableDefinition>,
    trees: &FnvHashMap<StringKey, Node>,
    base_definition_names: &FnvHashSet<StringKey>,
    key: StringKey,
) {
    if !visited.insert(key) {
        return;
    }

    let parents = match trees.get(&key) {
        None => {
            panic!("Fragment {:?} not found in IR.", key);
        }
        Some(node) => node.parents.clone(),
    };
    for parent in parents.into_iter() {
        add_related_nodes(visited, result, trees, base_definition_names, parent);
    }
    if !base_definition_names.contains(&key) {
        add_descendants(result, trees, key);
    }
}

// Recursively add all descendants of current node into the `result`
fn add_descendants(
    result: &mut FnvHashMap<StringKey, ExecutableDefinition>,
    trees: &FnvHashMap<StringKey, Node>,
    key: StringKey,
) {
    if result.contains_key(&key) {
        return;
    }
    match trees.get(&key) {
        Some(Node {
            ir: Some(def),
            children,
            ..
        }) => {
            result.insert(key, def.clone());
            for child in children {
                add_descendants(result, trees, *child);
            }
        }
        _ => {
            panic!("Fragment {:?} not found in IR.", key);
        }
    }
}
<commit_msg>Only traversing down from root when adding related IR dependencies<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use fnv::{FnvHashMap, FnvHashSet};

use graphql_ir::*;
use interner::StringKey;
use std::collections::hash_map::Entry;
use std::fmt;

struct Node {
    ir: Option<ExecutableDefinition>,
    parents: Vec<StringKey>,
    children: Vec<StringKey>,
}

impl fmt::Debug for Node {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_fmt(format_args!(
            "parents: {:?} / children: {:?}",
            self.parents, self.children
        ))
    }
}

pub fn get_reachable_ir(
    definitions: Vec<ExecutableDefinition>,
    base_definition_names: FnvHashSet<StringKey>,
    reachable_names: FnvHashSet<StringKey>,
) -> Vec<ExecutableDefinition> {
    if reachable_names.is_empty() {
        return vec![];
    }

    let trees = build_dependency_trees(definitions);

    let mut visited = FnvHashSet::default();
    let mut filtered_definitions = FnvHashMap::default();

    for key in reachable_names.into_iter() {
        if trees.contains_key(&key) {
            add_related_nodes(
                &mut visited,
                &mut filtered_definitions,
                &trees,
                &base_definition_names,
                key,
            );
        }
    }

    filtered_definitions
        .drain()
        .map(|(_, definition)| definition)
        .collect()
}

// Build a set of dependency trees that nodes are "doubly linked"
fn build_dependency_trees(definitions: Vec<ExecutableDefinition>) -> FnvHashMap<StringKey, Node> {
    let mut trees = FnvHashMap::with_capacity_and_hasher(definitions.len(), Default::default());

    for definition in definitions.into_iter() {
        let name = match &definition {
            ExecutableDefinition::Operation(operation) => operation.name.item,
            ExecutableDefinition::Fragment(fragment) => fragment.name.item,
        };

        // Visit the selections of the IR to build it's `children`
        let mut children = vec![];
        let selections = match &definition {
            ExecutableDefinition::Operation(operation) => &operation.selections,
            ExecutableDefinition::Fragment(fragment) => &fragment.selections,
        };
        visit_selections(&mut trees, &selections, name, &mut children);

        // Insert or update the representation of the IR in the dependency tree
        match trees.entry(name) {
            // Add a new node for current IR to the dependency tree
            Entry::Vacant(entry) => {
                entry.insert(Node {
                    ir: Some(definition),
                    parents: vec![],
                    children,
                });
            }

            // The node is already created when visiting selections of the IR, but it's `ir` and `children` haven't been set
            Entry::Occupied(mut entry) => {
                let node = entry.get_mut();
                if let Some(def) = &node.ir {
                    panic!(
                        "Duplicate definition: had {:?} and found another {:?}.",
                        def, node.ir
                    );
                }
                node.ir = Some(definition);
                node.children = children;
            }
        }
    }
    trees
}

// Visit the selections of current IR, set the `children` for the node representing the IR,
// and the `parents` for nodes representing the children IR
fn visit_selections(
    trees: &mut FnvHashMap<StringKey, Node>,
    selections: &[Selection],
    parent_name: StringKey,
    children: &mut Vec<StringKey>,
) {
    for selection in selections {
        match selection {
            Selection::FragmentSpread(node) => {
                let key = node.fragment.item;
                match trees.get_mut(&key) {
                    None => {
                        trees.insert(
                            key,
                            Node {
                                ir: None,
                                parents: vec![parent_name],
                                children: vec![],
                            },
                        );
                    }
                    Some(node) => {
                        node.parents.push(parent_name);
                    }
                }
                children.push(key);
            }
            Selection::LinkedField(node) => {
                visit_selections(trees, &node.selections, parent_name, children);
            }
            Selection::InlineFragment(node) => {
                visit_selections(trees, &node.selections, parent_name, children);
            }
            Selection::ScalarField(_) => {}
            Selection::Condition(node) => {
                visit_selections(trees, &node.selections, parent_name, children);
            }
        }
    }
}

// From `key` of changed definition, recusively traverse up the depenency tree, and add all related nodes (ancestors
// of changned definitions which are not from base definitions, and all of their desendants) into the `result`
fn add_related_nodes(
    visited: &mut FnvHashSet<StringKey>,
    result: &mut FnvHashMap<StringKey, ExecutableDefinition>,
    trees: &FnvHashMap<StringKey, Node>,
    base_definition_names: &FnvHashSet<StringKey>,
    key: StringKey,
) {
    if !visited.insert(key) {
        return;
    }

    let parents = match trees.get(&key) {
        None => {
            panic!("Fragment {:?} not found in IR.", key);
        }
        Some(node) => node.parents.clone(),
    };
    if parents.is_empty() {
        if !base_definition_names.contains(&key) {
            add_descendants(result, trees, key);
        }
    } else {
        for parent in parents.into_iter() {
            add_related_nodes(visited, result, trees, base_definition_names, parent);
        }
    }
}

// Recursively add all descendants of current node into the `result`
fn add_descendants(
    result: &mut FnvHashMap<StringKey, ExecutableDefinition>,
    trees: &FnvHashMap<StringKey, Node>,
    key: StringKey,
) {
    if result.contains_key(&key) {
        return;
    }
    match trees.get(&key) {
        Some(Node {
            ir: Some(def),
            children,
            ..
        }) => {
            result.insert(key, def.clone());
            for child in children {
                add_descendants(result, trees, *child);
            }
        }
        _ => {
            panic!("Fragment {:?} not found in IR.", key);
        }
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

//! An LSP-specific Compiler interface

use crate::lsp::{Connection, LSPBridgeMessage};

use relay_compiler::compiler_state::{CompilerState, ProjectName};
use relay_compiler::config::Config;
use relay_compiler::errors::Error as CompilerError;
use relay_compiler::FileSourceSubscription;
use relay_compiler::{build_schema, check_project, GraphQLAsts, Programs};
use schema::Schema;

use common::{DiagnosticsResult, PerfLogEvent, PerfLogger};
use interner::StringKey;

use crate::completion::GraphQLSourceCache;

use crate::error_reporting::{add_build_project_errors, add_syntax_errors};
use crate::state::ServerState;
use crate::text_documents::{
    on_did_change_text_document, on_did_close_text_document, on_did_open_text_document,
};

use crate::error::{LSPError, Result};

use common::ConsoleLogger;
use std::{collections::HashMap, sync::Arc};
use tokio::sync::mpsc::Receiver;

use tokio::select;

type SchemaMap = HashMap<ProjectName, Arc<Schema>>;

pub struct LSPCompiler<'schema, 'config> {
    lsp_rx: Receiver<LSPBridgeMessage>,
    schemas: &'schema SchemaMap,
    config: &'config Config,
    subscription: FileSourceSubscription<'config>,
    compiler_state: CompilerState,
    connection: Connection,
    synced_graphql_documents: GraphQLSourceCache,
    server_state: ServerState,
    project_programs: HashMap<StringKey, Programs>,
}

impl<'schema, 'config> LSPCompiler<'schema, 'config> {
    pub fn new(
        schemas: &'schema SchemaMap,
        config: &'config Config,
        subscription: FileSourceSubscription<'config>,
        compiler_state: CompilerState,
        lsp_rx: Receiver<LSPBridgeMessage>,
        connection: Connection,
    ) -> Self {
        let root_dir = config.root_dir.clone();
        let server_state = ServerState::new(root_dir);
        LSPCompiler {
            lsp_rx,
            config,
            schemas,
            subscription,
            compiler_state,
            connection,
            synced_graphql_documents: HashMap::new(),
            server_state,
            project_programs: HashMap::new(),
        }
    }

    async fn check_projects_and_report_errors(&mut self, event: &impl PerfLogEvent) {
        match self.check_projects(event).await {
            Ok(_) => {
                self.server_state.clear_diagnostics();
            }
            Err(err) => {
                self.server_state.clear_diagnostics();
                if let LSPError::CompilerError(err) = err {
                    match err {
                        CompilerError::DiagnosticsError { errors } => {
                            add_syntax_errors(errors, &mut self.server_state)
                        }
                        CompilerError::BuildProjectsErrors { errors } => {
                            add_build_project_errors(errors, &mut self.server_state)
                        }
                        // Ignore the rest of these errors for now
                        CompilerError::ConfigFileRead { .. } => {}
                        CompilerError::ConfigFileParse { .. } => {}
                        CompilerError::ConfigFileValidation { .. } => {}
                        CompilerError::ReadFileError { .. } => {}
                        CompilerError::WriteFileError { .. } => {}
                        CompilerError::SerializationError { .. } => {}
                        CompilerError::DeserializationError { .. } => {}
                        CompilerError::CanonicalizeRoot { .. } => {}
                        CompilerError::Watchman { .. } => {}
                        CompilerError::EmptyQueryResult => {}
                        CompilerError::FileRead { .. } => {}
                        CompilerError::Syntax { .. } => {}
                        CompilerError::JoinError { .. } => {}
                        CompilerError::PostArtifactsError { .. } => {}
                    }
                } else {
                    // TODO(brandondail) log here
                }
            }
        }
        self.server_state.commit_diagnostics(&self.connection);
    }

    pub async fn watch(&mut self) -> Result<()> {
        loop {
            select! {
                changes = self.subscription.next_change() => {
                    if let Some(file_source_changes) = changes? {
                        let incremental_check_event =
                        ConsoleLogger.create_event("incremental_check_event");
                    let incremental_check_time =
                        incremental_check_event.start("incremental_check_time");
                    let had_new_changes = self.compiler_state.merge_file_source_changes(
                        &self.config,
                        &file_source_changes,
                        &incremental_check_event,
                        &ConsoleLogger,
                        false,
                    )?;

                    if had_new_changes {
                        self.check_projects_and_report_errors(&incremental_check_event).await;
                    }

                    incremental_check_event.stop(incremental_check_time);
                    ConsoleLogger.complete_event(incremental_check_event);
                    // We probably don't want the messages queue to grow indefinitely
                    // and we need to flush then, as the check/build is completed
                    ConsoleLogger.flush();

                    }
                }
                message = self.lsp_rx.recv() => {
                    if let Some(message) = message {
                      self.on_lsp_bridge_message(message);
                    }
                }
            }
        }
    }

    fn on_lsp_bridge_message(&mut self, message: LSPBridgeMessage) {
        match message {
            // Completion request
            LSPBridgeMessage::CompletionRequest { .. } => {
                /* TODO: Re-enable auto-complete
                if let Some(completion_request) =
                    get_completion_request(params, &self.synced_graphql_documents)
                {
                    info!("completion_request {:#?}", self.project_programs.keys());
                    // TODO(brandondail) don't hardcode schema here
                    let project_key = "facebook-test".intern();
                    if let Some(schema) = self.schemas.get(&project_key) {
                        let programs = self.project_programs.get(&project_key);
                        info!("programs? {:?}", programs.is_some());
                        if let Some(items) =
                            completion_items_for_request(completion_request, schema, programs)
                        {
                            send_completion_response(items, request_id, &self.connection);
                        }
                    }
                }
                */
            }
            LSPBridgeMessage::DidOpenTextDocument(params) => {
                on_did_open_text_document(params, &mut self.synced_graphql_documents);
            }
            LSPBridgeMessage::DidChangeTextDocument(params) => {
                on_did_change_text_document(params, &mut self.synced_graphql_documents);
            }
            LSPBridgeMessage::DidCloseTextDocument(params) => {
                on_did_close_text_document(params, &mut self.synced_graphql_documents);
            }
        }
    }

    pub fn build_schemas(
        config: &Config,
        compiler_state: &CompilerState,
        setup_event: &impl PerfLogEvent,
    ) -> DiagnosticsResult<SchemaMap> {
        let timer = setup_event.start("build_schemas");
        let mut schemas = HashMap::new();
        for project_config in config.enabled_projects() {
            let schema = build_schema(compiler_state, project_config)?;
            schemas.insert(project_config.name, Arc::new(schema));
        }
        setup_event.stop(timer);
        Ok(schemas)
    }

    async fn check_projects(&mut self, setup_event: &impl PerfLogEvent) -> Result<()> {
        let graphql_asts = setup_event.time("parse_sources_time", || {
            GraphQLAsts::from_graphql_sources_map(
                &self.compiler_state.graphql_sources,
                &self.compiler_state.get_dirty_defintions(&self.config),
            )
        })?;
        let mut check_project_errors = vec![];
        let mut project_programs = HashMap::new();
        for project_config in self.config.enabled_projects() {
            if self
                .compiler_state
                .project_has_pending_changes(project_config.name)
            {
                let schema = Arc::clone(self.schemas.get(&project_config.name).unwrap());
                // TODO: consider running all projects in parallel
                let programs = check_project(
                    &self.config,
                    project_config,
                    &self.compiler_state,
                    &graphql_asts,
                    schema,
                    Arc::new(ConsoleLogger),
                )
                .map_err(|err| {
                    check_project_errors.push(err);
                })
                .ok();
                if let Some(programs) = programs {
                    project_programs.insert(project_config.name, programs);
                }
            }
        }

        if check_project_errors.is_empty() {
            self.project_programs = project_programs;
            Ok(())
        } else {
            Err(CompilerError::BuildProjectsErrors {
                errors: check_project_errors,
            }
            .into())
        }
    }
}
<commit_msg>fix tabbing<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

//! An LSP-specific Compiler interface

use crate::lsp::{Connection, LSPBridgeMessage};

use relay_compiler::compiler_state::{CompilerState, ProjectName};
use relay_compiler::config::Config;
use relay_compiler::errors::Error as CompilerError;
use relay_compiler::FileSourceSubscription;
use relay_compiler::{build_schema, check_project, GraphQLAsts, Programs};
use schema::Schema;

use common::{DiagnosticsResult, PerfLogEvent, PerfLogger};
use interner::StringKey;

use crate::completion::GraphQLSourceCache;

use crate::error_reporting::{add_build_project_errors, add_syntax_errors};
use crate::state::ServerState;
use crate::text_documents::{
    on_did_change_text_document, on_did_close_text_document, on_did_open_text_document,
};

use crate::error::{LSPError, Result};

use common::ConsoleLogger;
use std::{collections::HashMap, sync::Arc};
use tokio::sync::mpsc::Receiver;

use tokio::select;

type SchemaMap = HashMap<ProjectName, Arc<Schema>>;

pub struct LSPCompiler<'schema, 'config> {
    lsp_rx: Receiver<LSPBridgeMessage>,
    schemas: &'schema SchemaMap,
    config: &'config Config,
    subscription: FileSourceSubscription<'config>,
    compiler_state: CompilerState,
    connection: Connection,
    synced_graphql_documents: GraphQLSourceCache,
    server_state: ServerState,
    project_programs: HashMap<StringKey, Programs>,
}

impl<'schema, 'config> LSPCompiler<'schema, 'config> {
    pub fn new(
        schemas: &'schema SchemaMap,
        config: &'config Config,
        subscription: FileSourceSubscription<'config>,
        compiler_state: CompilerState,
        lsp_rx: Receiver<LSPBridgeMessage>,
        connection: Connection,
    ) -> Self {
        let root_dir = config.root_dir.clone();
        let server_state = ServerState::new(root_dir);
        LSPCompiler {
            lsp_rx,
            config,
            schemas,
            subscription,
            compiler_state,
            connection,
            synced_graphql_documents: HashMap::new(),
            server_state,
            project_programs: HashMap::new(),
        }
    }

    async fn check_projects_and_report_errors(&mut self, event: &impl PerfLogEvent) {
        match self.check_projects(event).await {
            Ok(_) => {
                self.server_state.clear_diagnostics();
            }
            Err(err) => {
                self.server_state.clear_diagnostics();
                if let LSPError::CompilerError(err) = err {
                    match err {
                        CompilerError::DiagnosticsError { errors } => {
                            add_syntax_errors(errors, &mut self.server_state)
                        }
                        CompilerError::BuildProjectsErrors { errors } => {
                            add_build_project_errors(errors, &mut self.server_state)
                        }
                        // Ignore the rest of these errors for now
                        CompilerError::ConfigFileRead { .. } => {}
                        CompilerError::ConfigFileParse { .. } => {}
                        CompilerError::ConfigFileValidation { .. } => {}
                        CompilerError::ReadFileError { .. } => {}
                        CompilerError::WriteFileError { .. } => {}
                        CompilerError::SerializationError { .. } => {}
                        CompilerError::DeserializationError { .. } => {}
                        CompilerError::CanonicalizeRoot { .. } => {}
                        CompilerError::Watchman { .. } => {}
                        CompilerError::EmptyQueryResult => {}
                        CompilerError::FileRead { .. } => {}
                        CompilerError::Syntax { .. } => {}
                        CompilerError::JoinError { .. } => {}
                        CompilerError::PostArtifactsError { .. } => {}
                    }
                } else {
                    // TODO(brandondail) log here
                }
            }
        }
        self.server_state.commit_diagnostics(&self.connection);
    }

    pub async fn watch(&mut self) -> Result<()> {
        loop {
            select! {
                changes = self.subscription.next_change() => {
                    if let Some(file_source_changes) = changes? {
                        let incremental_check_event =
                        ConsoleLogger.create_event("incremental_check_event");
                        let incremental_check_time =
                            incremental_check_event.start("incremental_check_time");
                        let had_new_changes = self.compiler_state.merge_file_source_changes(
                            &self.config,
                            &file_source_changes,
                            &incremental_check_event,
                            &ConsoleLogger,
                            false,
                        )?;

                        if had_new_changes {
                            self.check_projects_and_report_errors(&incremental_check_event).await;
                        }

                        incremental_check_event.stop(incremental_check_time);
                        ConsoleLogger.complete_event(incremental_check_event);
                        // We probably don't want the messages queue to grow indefinitely
                        // and we need to flush then, as the check/build is completed
                        ConsoleLogger.flush();

                    }
                }
                message = self.lsp_rx.recv() => {
                    if let Some(message) = message {
                      self.on_lsp_bridge_message(message);
                    }
                }
            }
        }
    }

    fn on_lsp_bridge_message(&mut self, message: LSPBridgeMessage) {
        match message {
            // Completion request
            LSPBridgeMessage::CompletionRequest { .. } => {
                /* TODO: Re-enable auto-complete
                if let Some(completion_request) =
                    get_completion_request(params, &self.synced_graphql_documents)
                {
                    info!("completion_request {:#?}", self.project_programs.keys());
                    // TODO(brandondail) don't hardcode schema here
                    let project_key = "facebook-test".intern();
                    if let Some(schema) = self.schemas.get(&project_key) {
                        let programs = self.project_programs.get(&project_key);
                        info!("programs? {:?}", programs.is_some());
                        if let Some(items) =
                            completion_items_for_request(completion_request, schema, programs)
                        {
                            send_completion_response(items, request_id, &self.connection);
                        }
                    }
                }
                */
            }
            LSPBridgeMessage::DidOpenTextDocument(params) => {
                on_did_open_text_document(params, &mut self.synced_graphql_documents);
            }
            LSPBridgeMessage::DidChangeTextDocument(params) => {
                on_did_change_text_document(params, &mut self.synced_graphql_documents);
            }
            LSPBridgeMessage::DidCloseTextDocument(params) => {
                on_did_close_text_document(params, &mut self.synced_graphql_documents);
            }
        }
    }

    pub fn build_schemas(
        config: &Config,
        compiler_state: &CompilerState,
        setup_event: &impl PerfLogEvent,
    ) -> DiagnosticsResult<SchemaMap> {
        let timer = setup_event.start("build_schemas");
        let mut schemas = HashMap::new();
        for project_config in config.enabled_projects() {
            let schema = build_schema(compiler_state, project_config)?;
            schemas.insert(project_config.name, Arc::new(schema));
        }
        setup_event.stop(timer);
        Ok(schemas)
    }

    async fn check_projects(&mut self, setup_event: &impl PerfLogEvent) -> Result<()> {
        let graphql_asts = setup_event.time("parse_sources_time", || {
            GraphQLAsts::from_graphql_sources_map(
                &self.compiler_state.graphql_sources,
                &self.compiler_state.get_dirty_defintions(&self.config),
            )
        })?;
        let mut check_project_errors = vec![];
        let mut project_programs = HashMap::new();
        for project_config in self.config.enabled_projects() {
            if self
                .compiler_state
                .project_has_pending_changes(project_config.name)
            {
                let schema = Arc::clone(self.schemas.get(&project_config.name).unwrap());
                // TODO: consider running all projects in parallel
                let programs = check_project(
                    &self.config,
                    project_config,
                    &self.compiler_state,
                    &graphql_asts,
                    schema,
                    Arc::new(ConsoleLogger),
                )
                .map_err(|err| {
                    check_project_errors.push(err);
                })
                .ok();
                if let Some(programs) = programs {
                    project_programs.insert(project_config.name, programs);
                }
            }
        }

        if check_project_errors.is_empty() {
            self.project_programs = project_programs;
            Ok(())
        } else {
            Err(CompilerError::BuildProjectsErrors {
                errors: check_project_errors,
            }
            .into())
        }
    }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

///
/// Implementation of the callbacks that the CSS selector engine uses to query the DOM.
///

use core::str::eq_slice;
use newcss::select::SelectHandler;
use script::dom::node::{AbstractNode, LayoutView};

pub struct NodeSelectHandler {
    node: AbstractNode<LayoutView>,
}

fn with_node_name<R>(node: AbstractNode<LayoutView>, f: &fn(&str) -> R) -> R {
    if !node.is_element() {
        fail!(~"attempting to style non-element node");
    }
    do node.with_imm_element |element_n| {
        f(element_n.tag_name)
    }
}

impl SelectHandler<AbstractNode<LayoutView>> for NodeSelectHandler {
    // FIXME(tkuehn): placeholder to get servo to compile
    fn node_has_class(&self, node: &AbstractNode<LayoutView>, s: &str) -> bool {
        true
    }

    // FIXME(tkuehn): placeholder to get servo to compile
    fn with_node_classes<R>(&self, node: &AbstractNode<LayoutView>, f: &fn(Option<&str>) -> R) -> R {
        f(None)
    }

    fn with_node_name<R>(&self, node: &AbstractNode<LayoutView>, f: &fn(&str) -> R) -> R {
        with_node_name(*node, f)
    }

    fn named_parent_node(&self, node: &AbstractNode<LayoutView>, name: &str)
                         -> Option<AbstractNode<LayoutView>> {
        match node.parent_node() {
            Some(parent) => {
                do with_node_name(parent) |node_name| {
                    if eq_slice(name, node_name) {
                        Some(parent)
                    } else {
                        None
                    }
                }
            }
            None => None
        }
    }

    fn parent_node(&self, node: &AbstractNode<LayoutView>) -> Option<AbstractNode<LayoutView>> {
        node.parent_node()
    }

    // TODO: Use a Bloom filter.
    fn named_ancestor_node(&self, node: &AbstractNode<LayoutView>, name: &str)
                           -> Option<AbstractNode<LayoutView>> {
        let mut node = *node;
        loop {
            let parent = node.parent_node();
            match parent {
                Some(parent) => {
                    let mut found = false;
                    do with_node_name(parent) |node_name| {
                        if eq_slice(name, node_name) {
                            found = true;
                        }
                    }
                    if found {
                        return Some(parent);
                    }
                    node = parent;
                }
                None => return None
            }
        }
    }

    fn node_is_root(&self, node: &AbstractNode<LayoutView>) -> bool {
        self.parent_node(node).is_none()
    }

    fn node_is_link(&self, node: &AbstractNode<LayoutView>) -> bool {
        if node.is_element() {
            do node.with_imm_element |element| {
                "a" == element.tag_name
            }
        } else {
            false
        }
    }

    fn with_node_classes<R>(&self, node: &AbstractNode<LayoutView>, f: &fn(Option<&str>) -> R) -> R {
        if !node.is_element() {
            fail!(~"attempting to style non-element node");
        }
        do node.with_imm_element() |element_n| {
            f(element_n.get_attr("class"))
        }
    }

    fn node_has_class(&self, node: &AbstractNode<LayoutView>, class: &str) -> bool {
        if !node.is_element() {
            fail!(~"attempting to style non-element node");
        }
        do node.with_imm_element |element_n| {
            match element_n.get_attr("class") {
                None => false,
                Some(existing_classes) => {
                    let mut ret = false;
                    for str::each_split_char(existing_classes, ' ') |s| {
                        if s == class {
                            ret = true;
                            break;
                        }
                    }
                    ret
                }
            }
        }
    }

    fn with_node_id<R>(&self, node: &AbstractNode<LayoutView>, f: &fn(Option<&str>) -> R) -> R {
        if !node.is_element() {
            fail!(~"attempting to style non-element node");
        }
        do node.with_imm_element() |element_n| {
            f(element_n.get_attr("id"))
        }
    }

    fn node_has_id(&self, node: &AbstractNode<LayoutView>, id: &str) -> bool {
        if !node.is_element() {
            fail!(~"attempting to style non-element node");
        }
        do node.with_imm_element |element_n| {
            match element_n.get_attr("id") {
                None => false,
                Some(existing_id) => id == existing_id
            }
        }
    }
}
<commit_msg>Fix merge fallout which was disabling all CSS classes.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

///
/// Implementation of the callbacks that the CSS selector engine uses to query the DOM.
///

use core::str::eq_slice;
use newcss::select::SelectHandler;
use script::dom::node::{AbstractNode, LayoutView};

pub struct NodeSelectHandler {
    node: AbstractNode<LayoutView>,
}

fn with_node_name<R>(node: AbstractNode<LayoutView>, f: &fn(&str) -> R) -> R {
    if !node.is_element() {
        fail!(~"attempting to style non-element node");
    }
    do node.with_imm_element |element_n| {
        f(element_n.tag_name)
    }
}

impl SelectHandler<AbstractNode<LayoutView>> for NodeSelectHandler {
    fn with_node_name<R>(&self, node: &AbstractNode<LayoutView>, f: &fn(&str) -> R) -> R {
        with_node_name(*node, f)
    }

    fn named_parent_node(&self, node: &AbstractNode<LayoutView>, name: &str)
                         -> Option<AbstractNode<LayoutView>> {
        match node.parent_node() {
            Some(parent) => {
                do with_node_name(parent) |node_name| {
                    if eq_slice(name, node_name) {
                        Some(parent)
                    } else {
                        None
                    }
                }
            }
            None => None
        }
    }

    fn parent_node(&self, node: &AbstractNode<LayoutView>) -> Option<AbstractNode<LayoutView>> {
        node.parent_node()
    }

    // TODO: Use a Bloom filter.
    fn named_ancestor_node(&self, node: &AbstractNode<LayoutView>, name: &str)
                           -> Option<AbstractNode<LayoutView>> {
        let mut node = *node;
        loop {
            let parent = node.parent_node();
            match parent {
                Some(parent) => {
                    let mut found = false;
                    do with_node_name(parent) |node_name| {
                        if eq_slice(name, node_name) {
                            found = true;
                        }
                    }
                    if found {
                        return Some(parent);
                    }
                    node = parent;
                }
                None => return None
            }
        }
    }

    fn node_is_root(&self, node: &AbstractNode<LayoutView>) -> bool {
        self.parent_node(node).is_none()
    }

    fn node_is_link(&self, node: &AbstractNode<LayoutView>) -> bool {
        if node.is_element() {
            do node.with_imm_element |element| {
                "a" == element.tag_name
            }
        } else {
            false
        }
    }

    fn with_node_classes<R>(&self, node: &AbstractNode<LayoutView>, f: &fn(Option<&str>) -> R) -> R {
        if !node.is_element() {
            fail!(~"attempting to style non-element node");
        }
        do node.with_imm_element() |element_n| {
            f(element_n.get_attr("class"))
        }
    }

    fn node_has_class(&self, node: &AbstractNode<LayoutView>, class: &str) -> bool {
        if !node.is_element() {
            fail!(~"attempting to style non-element node");
        }
        do node.with_imm_element |element_n| {
            match element_n.get_attr("class") {
                None => false,
                Some(existing_classes) => {
                    let mut ret = false;
                    for str::each_split_char(existing_classes, ' ') |s| {
                        if s == class {
                            ret = true;
                            break;
                        }
                    }
                    ret
                }
            }
        }
    }

    fn with_node_id<R>(&self, node: &AbstractNode<LayoutView>, f: &fn(Option<&str>) -> R) -> R {
        if !node.is_element() {
            fail!(~"attempting to style non-element node");
        }
        do node.with_imm_element() |element_n| {
            f(element_n.get_attr("id"))
        }
    }

    fn node_has_id(&self, node: &AbstractNode<LayoutView>, id: &str) -> bool {
        if !node.is_element() {
            fail!(~"attempting to style non-element node");
        }
        do node.with_imm_element |element_n| {
            match element_n.get_attr("id") {
                None => false,
                Some(existing_id) => id == existing_id
            }
        }
    }
}
<|endoftext|>