<commit_before>use std::io::{self, Write, Read, Cursor};
use std::net::{TcpStream, Shutdown};
use std::time::Duration;
use std::borrow::BorrowMut;
use byteorder::{BigEndian, WriteBytesExt, ReadBytesExt};
use enum_primitive::FromPrimitive;
use {Function, Reason, Result, ExceptionCode, Error, Coil, binary, Client};

const MODBUS_PROTOCOL_TCP: u16 = 0x0000;
const MODBUS_TCP_DEFAULT_PORT: u16 = 502;
const MODBUS_HEADER_SIZE: usize = 7;
const MODBUS_MAX_PACKET_SIZE: usize = 260;

/// Config structure for more control over the tcp socket settings
#[derive(Clone, Copy)]
pub struct Config {
    /// The TCP port to use for communication (Default: `502`)
    pub tcp_port: u16,
    /// Timeout when reading from the TCP socket (Default: `infinite`)
    pub tcp_read_timeout: Option<Duration>,
    /// Timeout when writing to the TCP socket (Default: `infinite`)
    pub tcp_write_timeout: Option<Duration>,
    /// The modbus Unit Identifier used in the modbus layer (Default: `1`)
    pub modbus_uid: u8,
}

impl Default for Config {
    fn default() -> Config {
        Config {
            tcp_port: MODBUS_TCP_DEFAULT_PORT,
            tcp_read_timeout: None,
            tcp_write_timeout: None,
            modbus_uid: 1,
        }
    }
}

#[derive(Debug, PartialEq)]
struct Header {
    tid: u16,
    pid: u16,
    len: u16,
    uid: u8,
}

impl Header {
    fn new(transport: &mut Transport, len: u16) -> Header {
        Header {
            tid: transport.new_tid(),
            pid: MODBUS_PROTOCOL_TCP,
            len: len - MODBUS_HEADER_SIZE as u16,
            uid: transport.uid,
        }
    }

    fn pack(&self) -> Result<Vec<u8>> {
        let mut buff = vec![];
        buff.write_u16::<BigEndian>(self.tid)?;
        buff.write_u16::<BigEndian>(self.pid)?;
        buff.write_u16::<BigEndian>(self.len)?;
        buff.write_u8(self.uid)?;
        Ok(buff)
    }

    fn unpack(buff: &[u8]) -> Result<Header> {
        let mut rdr = Cursor::new(buff);
        Ok(Header {
            tid: rdr.read_u16::<BigEndian>()?,
            pid: rdr.read_u16::<BigEndian>()?,
            len: rdr.read_u16::<BigEndian>()?,
            uid: rdr.read_u8()?
        })
    }
}

/// Context object which holds state for all modbus operations.
pub struct Transport {
    tid: u16,
    uid: u8,
    stream: TcpStream,
}

impl Transport {
    /// Create a new context context object and connect it to `addr` on modbus-tcp default
    /// port (502)
    pub fn new(addr: &str) -> io::Result<Transport> {
        Self::new_with_cfg(addr, Config::default())
    }

    /// Create a new context object and connect it to `addr` on port `port`
    pub fn new_with_cfg(addr: &str, cfg: Config) -> io::Result<Transport> {
        match TcpStream::connect((addr, cfg.tcp_port)) {
            Ok(s) => {
                s.set_read_timeout(cfg.tcp_read_timeout)?;
                s.set_write_timeout(cfg.tcp_write_timeout)?;
                s.set_nodelay(true)?;
                Ok(Transport {
                    tid: 0,
                    uid: cfg.modbus_uid,
                    stream: s,
                })
            }
            Err(e) => Err(e),
        }
    }

    // Create a new transaction Id, incrementing the previous one.
    // The Id is wrapping around if the Id reaches `u16::MAX`.
    fn new_tid(&mut self) -> u16 {
        self.tid = self.tid.wrapping_add(1);
        self.tid
    }

    fn read(self: &mut Self, fun: Function) -> Result<Vec<u8>> {
        let packed_size = |v: u16| {
            v / 8 +
            if v % 8 > 0 {
                1
            } else {
                0
            }
        };
        let (addr, count, expected_bytes) = match fun {
            Function::ReadCoils(a, c) |
            Function::ReadDiscreteInputs(a, c) => (a, c, packed_size(c) as usize),
            Function::ReadHoldingRegisters(a, c) |
            Function::ReadInputRegisters(a, c) => (a, c, 2 * c as usize),
            _ => return Err(Error::InvalidFunction),
        };

        if count < 1 {
            return Err(Error::InvalidData(Reason::RecvBufferEmpty));
        }

        if count as usize > MODBUS_MAX_PACKET_SIZE {
            return Err(Error::InvalidData(Reason::UnexpectedReplySize));
        }

        let header = Header::new(self, MODBUS_HEADER_SIZE as u16 + 6u16);
        let mut buff = header.pack()?;
        buff.write_u8(fun.code())?;
        buff.write_u16::<BigEndian>(addr)?;
        buff.write_u16::<BigEndian>(count)?;

        match self.stream.write_all(&buff) {
            Ok(_s) => {
                let mut reply = vec![0; MODBUS_HEADER_SIZE + expected_bytes + 2];
                match self.stream.read(&mut reply) {
                    Ok(_s) => {
                        let resp_hd = Header::unpack(&reply[..MODBUS_HEADER_SIZE])?;
                        Transport::validate_response_header(&header, &resp_hd)?;
                        Transport::validate_response_code(&buff, &reply)?;
                        Transport::get_reply_data(&reply, expected_bytes)
                    }
                    Err(e) => Err(Error::Io(e)),
                }
            }
            Err(e) => Err(Error::Io(e)),
        }
    }

    fn validate_response_header(req: &Header, resp: &Header) -> Result<()> {
        if req.tid != resp.tid || resp.pid != MODBUS_PROTOCOL_TCP {
            Err(Error::InvalidResponse)
        } else {
            Ok(())
        }
    }

    fn validate_response_code(req: &[u8], resp: &[u8]) -> Result<()> {
        if req[7] + 0x80 == resp[7] {
            match ExceptionCode::from_u8(resp[8]) {
                Some(code) => Err(Error::Exception(code)),
                None => Err(Error::InvalidResponse),
            }
        } else if req[7] == resp[7] {
            Ok(())
        } else {
            Err(Error::InvalidResponse)
        }
    }

    fn get_reply_data(reply: &[u8], expected_bytes: usize) -> Result<Vec<u8>> {
        if reply[8] as usize != expected_bytes ||
           reply.len() != MODBUS_HEADER_SIZE + expected_bytes + 2 {
            Err(Error::InvalidData(Reason::UnexpectedReplySize))
        } else {
            let mut d = Vec::new();
            d.extend_from_slice(&reply[MODBUS_HEADER_SIZE + 2..]);
            Ok(d)
        }
    }

    fn write_single(self: &mut Self, fun: Function) -> Result<()> {
        let (addr, value) = match fun {
            Function::WriteSingleCoil(a, v) |
            Function::WriteSingleRegister(a, v) => (a, v),
            _ => return Err(Error::InvalidFunction),
        };

        let mut buff = vec![0; MODBUS_HEADER_SIZE];  // Header gets filled in later
        buff.write_u8(fun.code())?;
        buff.write_u16::<BigEndian>(addr)?;
        buff.write_u16::<BigEndian>(value)?;
        self.write(&mut buff)
    }

    fn write_multiple(self: &mut Self, fun: Function) -> Result<()> {
        let (addr, quantity, values) = match fun {
            Function::WriteMultipleCoils(a, q, v) |
            Function::WriteMultipleRegisters(a, q, v) => (a, q, v),
            _ => return Err(Error::InvalidFunction),
        };

        let mut buff = vec![0; MODBUS_HEADER_SIZE];  // Header gets filled in later
        buff.write_u8(fun.code())?;
        buff.write_u16::<BigEndian>(addr)?;
        buff.write_u16::<BigEndian>(quantity)?;
        buff.write_u8(values.len() as u8)?;
        for v in values {
            buff.write_u8(*v)?;
        }
        self.write(&mut buff)
    }

    fn write(self: &mut Self, buff: &mut [u8]) -> Result<()> {
        if buff.len() < 1 {
            return Err(Error::InvalidData(Reason::SendBufferEmpty));
        }

        if buff.len() > MODBUS_MAX_PACKET_SIZE {
            return Err(Error::InvalidData(Reason::SendBufferTooBig));
        }

        let header = Header::new(self, buff.len() as u16 + 1u16);
        let head_buff = header.pack()?;
        {
            let mut start = Cursor::new(buff.borrow_mut());
            start.write_all(&head_buff)?;
        }
        match self.stream.write_all(buff) {
            Ok(_s) => {
                let reply = &mut [0; 12];
                match self.stream.read(reply) {
                    Ok(_s) => {
                        let resp_hd = Header::unpack(reply)?;
                        Transport::validate_response_header(&header, &resp_hd)?;
                        Transport::validate_response_code(buff, reply)
                    }
                    Err(e) => Err(Error::Io(e)),
                }
            }
            Err(e) => Err(Error::Io(e)),
        }
    }

    pub fn close(self: &mut Self) -> Result<()> {
        self.stream.shutdown(Shutdown::Both).map_err(Error::Io)
    }
}

impl Client for Transport {
    /// Read `count` bits starting at address `addr`.
    fn read_coils(self: &mut Self, addr: u16, count: u16) -> Result<Vec<Coil>> {
        let bytes = self.read(Function::ReadCoils(addr, count))?;
        Ok(binary::unpack_bits(&bytes, count))
    }

    /// Read `count` input bits starting at address `addr`.
    fn read_discrete_inputs(self: &mut Self, addr: u16, count: u16) -> Result<Vec<Coil>> {
        let bytes = self.read(Function::ReadDiscreteInputs(addr, count))?;
        Ok(binary::unpack_bits(&bytes, count))
    }

    /// Read `count` 16bit registers starting at address `addr`.
    fn read_holding_registers(self: &mut Self, addr: u16, count: u16) -> Result<Vec<u16>> {
        let bytes = self.read(Function::ReadHoldingRegisters(addr, count))?;
        binary::pack_bytes(&bytes[..])
    }

    /// Read `count` 16bit input registers starting at address `addr`.
    fn read_input_registers(self: &mut Self, addr: u16, count: u16) -> Result<Vec<u16>> {
        let bytes = self.read(Function::ReadInputRegisters(addr, count))?;
        binary::pack_bytes(&bytes[..])
    }

    /// Write a single coil (bit) to address `addr`.
    fn write_single_coil(self: &mut Self, addr: u16, value: Coil) -> Result<()> {
        self.write_single(Function::WriteSingleCoil(addr, value.code()))
    }

    /// Write a single 16bit register to address `addr`.
    fn write_single_register(self: &mut Self, addr: u16, value: u16) -> Result<()> {
        self.write_single(Function::WriteSingleRegister(addr, value))
    }

    /// Write a multiple coils (bits) starting at address `addr`.
    fn write_multiple_coils(self: &mut Self, addr: u16, values: &[Coil]) -> Result<()> {
        let bytes = binary::pack_bits(values);
        self.write_multiple(Function::WriteMultipleCoils(addr, values.len() as u16, &bytes))
    }

    /// Write a multiple 16bit registers starting at address `addr`.
    fn write_multiple_registers(self: &mut Self, addr: u16, values: &[u16]) -> Result<()> {
        let bytes = binary::unpack_bytes(values);
        self.write_multiple(Function::WriteMultipleRegisters(addr, values.len() as u16, &bytes))
    }
}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn serialize_header(){
        let header = Header{tid: 12816, pid: 3930, len: 99, uid: 68 };
        let serialized = header.pack().unwrap();
        let deserialized = Header::unpack(&vec![50, 16, 15, 90, 0, 99, 68]).unwrap();
        let re_deserialized = Header::unpack(&serialized).unwrap();
        assert_eq!(serialized, vec![50, 16, 15, 90, 0, 99, 68]);
        assert_eq!(deserialized, header);
        assert_eq!(re_deserialized, header);
    }
}
<commit_msg>Apply some clippy suggestions<commit_after>use std::io::{self, Write, Read, Cursor};
use std::net::{TcpStream, Shutdown};
use std::time::Duration;
use std::borrow::BorrowMut;
use byteorder::{BigEndian, WriteBytesExt, ReadBytesExt};
use enum_primitive::FromPrimitive;
use {Function, Reason, Result, ExceptionCode, Error, Coil, binary, Client};

const MODBUS_PROTOCOL_TCP: u16 = 0x0000;
const MODBUS_TCP_DEFAULT_PORT: u16 = 502;
const MODBUS_HEADER_SIZE: usize = 7;
const MODBUS_MAX_PACKET_SIZE: usize = 260;

/// Config structure for more control over the tcp socket settings
#[derive(Clone, Copy)]
pub struct Config {
    /// The TCP port to use for communication (Default: `502`)
    pub tcp_port: u16,
    /// Timeout when reading from the TCP socket (Default: `infinite`)
    pub tcp_read_timeout: Option<Duration>,
    /// Timeout when writing to the TCP socket (Default: `infinite`)
    pub tcp_write_timeout: Option<Duration>,
    /// The modbus Unit Identifier used in the modbus layer (Default: `1`)
    pub modbus_uid: u8,
}

impl Default for Config {
    fn default() -> Config {
        Config {
            tcp_port: MODBUS_TCP_DEFAULT_PORT,
            tcp_read_timeout: None,
            tcp_write_timeout: None,
            modbus_uid: 1,
        }
    }
}

#[derive(Debug, PartialEq)]
struct Header {
    tid: u16,
    pid: u16,
    len: u16,
    uid: u8,
}

impl Header {
    fn new(transport: &mut Transport, len: u16) -> Header {
        Header {
            tid: transport.new_tid(),
            pid: MODBUS_PROTOCOL_TCP,
            len: len - MODBUS_HEADER_SIZE as u16,
            uid: transport.uid,
        }
    }

    fn pack(&self) -> Result<Vec<u8>> {
        let mut buff = vec![];
        buff.write_u16::<BigEndian>(self.tid)?;
        buff.write_u16::<BigEndian>(self.pid)?;
        buff.write_u16::<BigEndian>(self.len)?;
        buff.write_u8(self.uid)?;
        Ok(buff)
    }

    fn unpack(buff: &[u8]) -> Result<Header> {
        let mut rdr = Cursor::new(buff);
        Ok(Header {
            tid: rdr.read_u16::<BigEndian>()?,
            pid: rdr.read_u16::<BigEndian>()?,
            len: rdr.read_u16::<BigEndian>()?,
            uid: rdr.read_u8()?
        })
    }
}

/// Context object which holds state for all modbus operations.
pub struct Transport {
    tid: u16,
    uid: u8,
    stream: TcpStream,
}

impl Transport {
    /// Create a new context context object and connect it to `addr` on modbus-tcp default
    /// port (502)
    pub fn new(addr: &str) -> io::Result<Transport> {
        Self::new_with_cfg(addr, Config::default())
    }

    /// Create a new context object and connect it to `addr` on port `port`
    pub fn new_with_cfg(addr: &str, cfg: Config) -> io::Result<Transport> {
        match TcpStream::connect((addr, cfg.tcp_port)) {
            Ok(s) => {
                s.set_read_timeout(cfg.tcp_read_timeout)?;
                s.set_write_timeout(cfg.tcp_write_timeout)?;
                s.set_nodelay(true)?;
                Ok(Transport {
                    tid: 0,
                    uid: cfg.modbus_uid,
                    stream: s,
                })
            }
            Err(e) => Err(e),
        }
    }

    // Create a new transaction Id, incrementing the previous one.
    // The Id is wrapping around if the Id reaches `u16::MAX`.
    fn new_tid(&mut self) -> u16 {
        self.tid = self.tid.wrapping_add(1);
        self.tid
    }

    fn read(self: &mut Self, fun: &Function) -> Result<Vec<u8>> {
        let packed_size = |v: u16| {
            v / 8 +
            if v % 8 > 0 {
                1
            } else {
                0
            }
        };
        let (addr, count, expected_bytes) = match *fun {
            Function::ReadCoils(a, c) |
            Function::ReadDiscreteInputs(a, c) => (a, c, packed_size(c) as usize),
            Function::ReadHoldingRegisters(a, c) |
            Function::ReadInputRegisters(a, c) => (a, c, 2 * c as usize),
            _ => return Err(Error::InvalidFunction),
        };

        if count < 1 {
            return Err(Error::InvalidData(Reason::RecvBufferEmpty));
        }

        if count as usize > MODBUS_MAX_PACKET_SIZE {
            return Err(Error::InvalidData(Reason::UnexpectedReplySize));
        }

        let header = Header::new(self, MODBUS_HEADER_SIZE as u16 + 6u16);
        let mut buff = header.pack()?;
        buff.write_u8(fun.code())?;
        buff.write_u16::<BigEndian>(addr)?;
        buff.write_u16::<BigEndian>(count)?;

        match self.stream.write_all(&buff) {
            Ok(_s) => {
                let mut reply = vec![0; MODBUS_HEADER_SIZE + expected_bytes + 2];
                match self.stream.read(&mut reply) {
                    Ok(_s) => {
                        let resp_hd = Header::unpack(&reply[..MODBUS_HEADER_SIZE])?;
                        Transport::validate_response_header(&header, &resp_hd)?;
                        Transport::validate_response_code(&buff, &reply)?;
                        Transport::get_reply_data(&reply, expected_bytes)
                    }
                    Err(e) => Err(Error::Io(e)),
                }
            }
            Err(e) => Err(Error::Io(e)),
        }
    }

    fn validate_response_header(req: &Header, resp: &Header) -> Result<()> {
        if req.tid != resp.tid || resp.pid != MODBUS_PROTOCOL_TCP {
            Err(Error::InvalidResponse)
        } else {
            Ok(())
        }
    }

    fn validate_response_code(req: &[u8], resp: &[u8]) -> Result<()> {
        if req[7] + 0x80 == resp[7] {
            match ExceptionCode::from_u8(resp[8]) {
                Some(code) => Err(Error::Exception(code)),
                None => Err(Error::InvalidResponse),
            }
        } else if req[7] == resp[7] {
            Ok(())
        } else {
            Err(Error::InvalidResponse)
        }
    }

    fn get_reply_data(reply: &[u8], expected_bytes: usize) -> Result<Vec<u8>> {
        if reply[8] as usize != expected_bytes ||
           reply.len() != MODBUS_HEADER_SIZE + expected_bytes + 2 {
            Err(Error::InvalidData(Reason::UnexpectedReplySize))
        } else {
            let mut d = Vec::new();
            d.extend_from_slice(&reply[MODBUS_HEADER_SIZE + 2..]);
            Ok(d)
        }
    }

    fn write_single(self: &mut Self, fun: &Function) -> Result<()> {
        let (addr, value) = match *fun {
            Function::WriteSingleCoil(a, v) |
            Function::WriteSingleRegister(a, v) => (a, v),
            _ => return Err(Error::InvalidFunction),
        };

        let mut buff = vec![0; MODBUS_HEADER_SIZE];  // Header gets filled in later
        buff.write_u8(fun.code())?;
        buff.write_u16::<BigEndian>(addr)?;
        buff.write_u16::<BigEndian>(value)?;
        self.write(&mut buff)
    }

    fn write_multiple(self: &mut Self, fun: &Function) -> Result<()> {
        let (addr, quantity, values) = match *fun {
            Function::WriteMultipleCoils(a, q, v) |
            Function::WriteMultipleRegisters(a, q, v) => (a, q, v),
            _ => return Err(Error::InvalidFunction),
        };

        let mut buff = vec![0; MODBUS_HEADER_SIZE];  // Header gets filled in later
        buff.write_u8(fun.code())?;
        buff.write_u16::<BigEndian>(addr)?;
        buff.write_u16::<BigEndian>(quantity)?;
        buff.write_u8(values.len() as u8)?;
        for v in values {
            buff.write_u8(*v)?;
        }
        self.write(&mut buff)
    }

    fn write(self: &mut Self, buff: &mut [u8]) -> Result<()> {
        if buff.len() < 1 {
            return Err(Error::InvalidData(Reason::SendBufferEmpty));
        }

        if buff.len() > MODBUS_MAX_PACKET_SIZE {
            return Err(Error::InvalidData(Reason::SendBufferTooBig));
        }

        let header = Header::new(self, buff.len() as u16 + 1u16);
        let head_buff = header.pack()?;
        {
            let mut start = Cursor::new(buff.borrow_mut());
            start.write_all(&head_buff)?;
        }
        match self.stream.write_all(buff) {
            Ok(_s) => {
                let reply = &mut [0; 12];
                match self.stream.read(reply) {
                    Ok(_s) => {
                        let resp_hd = Header::unpack(reply)?;
                        Transport::validate_response_header(&header, &resp_hd)?;
                        Transport::validate_response_code(buff, reply)
                    }
                    Err(e) => Err(Error::Io(e)),
                }
            }
            Err(e) => Err(Error::Io(e)),
        }
    }

    pub fn close(self: &mut Self) -> Result<()> {
        self.stream.shutdown(Shutdown::Both).map_err(Error::Io)
    }
}

impl Client for Transport {
    /// Read `count` bits starting at address `addr`.
    fn read_coils(self: &mut Self, addr: u16, count: u16) -> Result<Vec<Coil>> {
        let bytes = self.read(&Function::ReadCoils(addr, count))?;
        Ok(binary::unpack_bits(&bytes, count))
    }

    /// Read `count` input bits starting at address `addr`.
    fn read_discrete_inputs(self: &mut Self, addr: u16, count: u16) -> Result<Vec<Coil>> {
        let bytes = self.read(&Function::ReadDiscreteInputs(addr, count))?;
        Ok(binary::unpack_bits(&bytes, count))
    }

    /// Read `count` 16bit registers starting at address `addr`.
    fn read_holding_registers(self: &mut Self, addr: u16, count: u16) -> Result<Vec<u16>> {
        let bytes = self.read(&Function::ReadHoldingRegisters(addr, count))?;
        binary::pack_bytes(&bytes[..])
    }

    /// Read `count` 16bit input registers starting at address `addr`.
    fn read_input_registers(self: &mut Self, addr: u16, count: u16) -> Result<Vec<u16>> {
        let bytes = self.read(&Function::ReadInputRegisters(addr, count))?;
        binary::pack_bytes(&bytes[..])
    }

    /// Write a single coil (bit) to address `addr`.
    fn write_single_coil(self: &mut Self, addr: u16, value: Coil) -> Result<()> {
        self.write_single(&Function::WriteSingleCoil(addr, value.code()))
    }

    /// Write a single 16bit register to address `addr`.
    fn write_single_register(self: &mut Self, addr: u16, value: u16) -> Result<()> {
        self.write_single(&Function::WriteSingleRegister(addr, value))
    }

    /// Write a multiple coils (bits) starting at address `addr`.
    fn write_multiple_coils(self: &mut Self, addr: u16, values: &[Coil]) -> Result<()> {
        let bytes = binary::pack_bits(values);
        self.write_multiple(&Function::WriteMultipleCoils(addr, values.len() as u16, &bytes))
    }

    /// Write a multiple 16bit registers starting at address `addr`.
    fn write_multiple_registers(self: &mut Self, addr: u16, values: &[u16]) -> Result<()> {
        let bytes = binary::unpack_bytes(values);
        self.write_multiple(&Function::WriteMultipleRegisters(addr, values.len() as u16, &bytes))
    }
}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn serialize_header(){
        let header = Header{tid: 12816, pid: 3930, len: 99, uid: 68 };
        let serialized = header.pack().unwrap();
        let deserialized = Header::unpack(&vec![50, 16, 15, 90, 0, 99, 68]).unwrap();
        let re_deserialized = Header::unpack(&serialized).unwrap();
        assert_eq!(serialized, vec![50, 16, 15, 90, 0, 99, 68]);
        assert_eq!(deserialized, header);
        assert_eq!(re_deserialized, header);
    }
}
<|endoftext|>


<commit_before>#[cfg(windows)]
use std::sync::Mutex;

use super::*;

#[cfg(unix)]
use std::os::unix::fs::FileExt;

#[cfg(windows)]
use std::os::windows::fs::FileExt;

/// Multithreaded IO support for Files
pub trait Pio {
    /// Read from a specific offset without changing
    /// the underlying file offset.
    fn pread_exact(
        &self,
        mut buf: &mut [u8],
        mut offset: LogID,
    ) -> io::Result<()>;

    /// Write to a specific offset without changing
    /// the underlying file offset.
    fn pwrite_all(&self, mut buf: &[u8], mut offset: LogID) -> io::Result<()>;
}

// On systems that support pread/pwrite, use them underneath.
#[cfg(unix)]
impl Pio for std::fs::File {
    fn pread_exact(
        &self,
        mut buf: &mut [u8],
        mut offset: LogID,
    ) -> io::Result<()> {
        while !buf.is_empty() {
            match self.read_at(buf, offset) {
                Ok(0) => break,
                Ok(n) => {
                    offset += n as LogID;
                    let tmp = buf;
                    buf = &mut tmp[n..];
                }
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
                Err(e) => return Err(e),
            }
        }
        if !buf.is_empty() {
            Err(io::Error::new(
                io::ErrorKind::UnexpectedEof,
                "failed to fill whole buffer",
            ))
        } else {
            Ok(())
        }
    }

    fn pwrite_all(&self, mut buf: &[u8], mut offset: LogID) -> io::Result<()> {
        while !buf.is_empty() {
            match self.write_at(buf, offset) {
                Ok(0) => {
                    return Err(io::Error::new(
                        io::ErrorKind::WriteZero,
                        "failed to write whole buffer",
                    ))
                }
                Ok(n) => {
                    offset += n as LogID;
                    buf = &buf[n..]
                }
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
                Err(e) => return Err(e),
            }
        }
        Ok(())
    }
}

// HACK HACK HACK get this working with real parallel IO
#[cfg(windows)]
lazy_static! {
    pub static ref GLOBAL_FILE_LOCK: Mutex<()> = Mutex::new(());
}

#[cfg(windows)]
impl Pio for std::fs::File {
    fn pread_exact(
        &self,
        mut buf: &mut [u8],
        mut offset: LogID,
    ) -> io::Result<()> {
        // HACK HACK HACK get this working with real parallel IO
        let _lock = GLOBAL_FILE_LOCK.lock().unwrap();

        while !buf.is_empty() {
            match self.seek_read(buf, offset) {
                Ok(0) => break,
                Ok(n) => {
                    offset += n as LogID;
                    let tmp = buf;
                    buf = &mut tmp[n..];
                }
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
                Err(e) => return Err(e),
            }
        }
        if !buf.is_empty() {
            Err(io::Error::new(
                io::ErrorKind::UnexpectedEof,
                "failed to fill whole buffer",
            ))
        } else {
            Ok(())
        }
    }

    fn pwrite_all(&self, mut buf: &[u8], mut offset: LogID) -> io::Result<()> {
        // HACK HACK HACK get this working with real parallel IO
        let _lock = GLOBAL_FILE_LOCK.lock().unwrap();

        while !buf.is_empty() {
            match self.seek_write(buf, offset) {
                Ok(0) => {
                    return Err(io::Error::new(
                        io::ErrorKind::WriteZero,
                        "failed to write whole buffer",
                    ))
                }
                Ok(n) => {
                    offset += n as LogID;
                    buf = &buf[n..]
                }
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
                Err(e) => return Err(e),
            }
        }
        Ok(())
    }
}
<commit_msg>remove mut pattern from Pio trait def<commit_after>#[cfg(windows)]
use std::sync::Mutex;

use super::*;

#[cfg(unix)]
use std::os::unix::fs::FileExt;

#[cfg(windows)]
use std::os::windows::fs::FileExt;

/// Multithreaded IO support for Files
pub trait Pio {
    /// Read from a specific offset without changing
    /// the underlying file offset.
    fn pread_exact(&self, to_buf: &mut [u8], offset: LogID) -> io::Result<()>;

    /// Write to a specific offset without changing
    /// the underlying file offset.
    fn pwrite_all(&self, from_buf: &[u8], offset: LogID) -> io::Result<()>;
}

// On systems that support pread/pwrite, use them underneath.
#[cfg(unix)]
impl Pio for std::fs::File {
    fn pread_exact(
        &self,
        mut buf: &mut [u8],
        mut offset: LogID,
    ) -> io::Result<()> {
        while !buf.is_empty() {
            match self.read_at(buf, offset) {
                Ok(0) => break,
                Ok(n) => {
                    offset += n as LogID;
                    let tmp = buf;
                    buf = &mut tmp[n..];
                }
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
                Err(e) => return Err(e),
            }
        }
        if !buf.is_empty() {
            Err(io::Error::new(
                io::ErrorKind::UnexpectedEof,
                "failed to fill whole buffer",
            ))
        } else {
            Ok(())
        }
    }

    fn pwrite_all(&self, mut buf: &[u8], mut offset: LogID) -> io::Result<()> {
        while !buf.is_empty() {
            match self.write_at(buf, offset) {
                Ok(0) => {
                    return Err(io::Error::new(
                        io::ErrorKind::WriteZero,
                        "failed to write whole buffer",
                    ))
                }
                Ok(n) => {
                    offset += n as LogID;
                    buf = &buf[n..]
                }
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
                Err(e) => return Err(e),
            }
        }
        Ok(())
    }
}

// HACK HACK HACK get this working with real parallel IO
#[cfg(windows)]
lazy_static! {
    pub static ref GLOBAL_FILE_LOCK: Mutex<()> = Mutex::new(());
}

#[cfg(windows)]
impl Pio for std::fs::File {
    fn pread_exact(
        &self,
        mut buf: &mut [u8],
        mut offset: LogID,
    ) -> io::Result<()> {
        // HACK HACK HACK get this working with real parallel IO
        let _lock = GLOBAL_FILE_LOCK.lock().unwrap();

        while !buf.is_empty() {
            match self.seek_read(buf, offset) {
                Ok(0) => break,
                Ok(n) => {
                    offset += n as LogID;
                    let tmp = buf;
                    buf = &mut tmp[n..];
                }
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
                Err(e) => return Err(e),
            }
        }
        if !buf.is_empty() {
            Err(io::Error::new(
                io::ErrorKind::UnexpectedEof,
                "failed to fill whole buffer",
            ))
        } else {
            Ok(())
        }
    }

    fn pwrite_all(&self, mut buf: &[u8], mut offset: LogID) -> io::Result<()> {
        // HACK HACK HACK get this working with real parallel IO
        let _lock = GLOBAL_FILE_LOCK.lock().unwrap();

        while !buf.is_empty() {
            match self.seek_write(buf, offset) {
                Ok(0) => {
                    return Err(io::Error::new(
                        io::ErrorKind::WriteZero,
                        "failed to write whole buffer",
                    ))
                }
                Ok(n) => {
                    offset += n as LogID;
                    buf = &buf[n..]
                }
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
                Err(e) => return Err(e),
            }
        }
        Ok(())
    }
}
<|endoftext|>


<commit_before>//! proc-macro tests

#[macro_use]
mod utils;
use utils::*;

use expect_test::expect;

#[test]
fn test_derive_empty() {
    assert_expand("DeriveEmpty", r#"struct S;"#, expect![[r#"SUBTREE $"#]]);
}

#[test]
fn test_derive_error() {
    assert_expand(
        "DeriveError",
        r#"struct S;"#,
        expect![[r##"
            SUBTREE $
              IDENT   compile_error 4294967295
              PUNCH   ! [joint] 4294967295
              SUBTREE () 4294967295
                LITERAL "#[derive(DeriveError)] struct S ;" 4294967295
              PUNCH   ; [alone] 4294967295"##]],
    );
}

#[test]
fn test_fn_like_macro_noop() {
    assert_expand(
        "fn_like_noop",
        r#"ident, 0, 1, []"#,
        expect![[r#"
            SUBTREE $
              IDENT   ident 4294967295
              PUNCH   , [alone] 4294967295
              LITERAL 0 4294967295
              PUNCH   , [alone] 4294967295
              LITERAL 1 4294967295
              PUNCH   , [alone] 4294967295
              SUBTREE [] 4294967295"#]],
    );
}

#[test]
fn test_fn_like_macro_clone_ident_subtree() {
    assert_expand(
        "fn_like_clone_tokens",
        r#"ident, []"#,
        expect![[r#"
            SUBTREE $
              IDENT   ident 4294967295
              PUNCH   , [alone] 4294967295
              SUBTREE [] 4294967295"#]],
    );
}

#[test]
fn test_fn_like_macro_clone_raw_ident() {
    assert_expand(
        "fn_like_clone_tokens",
        "r#async",
        expect![[r#"
            SUBTREE $
              IDENT   async 4294967295"#]],
    );
}

#[test]
fn test_fn_like_mk_literals() {
    assert_expand(
        "fn_like_mk_literals",
        r#""#,
        expect![[r#"
            SUBTREE $
              LITERAL b"byte_string" 4294967295
              LITERAL 'c' 4294967295
              LITERAL "string" 4294967295
              LITERAL 3.14f64 4294967295
              LITERAL 3.14 4294967295
              LITERAL 123i64 4294967295
              LITERAL 123 4294967295"#]],
    );
}

#[test]
fn test_fn_like_mk_idents() {
    // FIXME: this test is wrong: raw should be 'r#raw' but ABIs 1.64 and below
    // simply ignore `is_raw` when implementing the `Ident` interface.
    assert_expand(
        "fn_like_mk_idents",
        r#""#,
        expect![[r#"
            SUBTREE $
              IDENT   standard 4294967295
              IDENT   raw 4294967295"#]],
    );
}

#[test]
fn test_fn_like_macro_clone_literals() {
    assert_expand(
        "fn_like_clone_tokens",
        r#"1u16, 2_u32, -4i64, 3.14f32, "hello bridge""#,
        expect![[r#"
            SUBTREE $
              LITERAL 1u16 4294967295
              PUNCH   , [alone] 4294967295
              LITERAL 2_u32 4294967295
              PUNCH   , [alone] 4294967295
              PUNCH   - [joint] 4294967295
              LITERAL 4i64 4294967295
              PUNCH   , [alone] 4294967295
              LITERAL 3.14f32 4294967295
              PUNCH   , [alone] 4294967295
              LITERAL "hello bridge" 4294967295"#]],
    );
}

#[test]
fn test_attr_macro() {
    // Corresponds to
    //    #[proc_macro_test::attr_error(some arguments)]
    //    mod m {}
    assert_expand_attr(
        "attr_error",
        r#"mod m {}"#,
        r#"some arguments"#,
        expect![[r##"
            SUBTREE $
              IDENT   compile_error 4294967295
              PUNCH   ! [joint] 4294967295
              SUBTREE () 4294967295
                LITERAL "#[attr_error(some arguments)] mod m {}" 4294967295
              PUNCH   ; [alone] 4294967295"##]],
    );
}

/// Tests that we find and classify all proc macros correctly.
#[test]
fn list_test_macros() {
    let res = list().join("\n");

    expect![[r#"
        fn_like_noop [FuncLike]
        fn_like_panic [FuncLike]
        fn_like_error [FuncLike]
        fn_like_clone_tokens [FuncLike]
        fn_like_mk_literals [FuncLike]
        fn_like_mk_idents [FuncLike]
        attr_noop [Attr]
        attr_panic [Attr]
        attr_error [Attr]
        DeriveEmpty [CustomDerive]
        DerivePanic [CustomDerive]
        DeriveError [CustomDerive]"#]]
    .assert_eq(&res);
}
<commit_msg>Auto merge of #13625 - lowr:fix/proc-macro-src-test, r=lnicola<commit_after>//! proc-macro tests

#[macro_use]
mod utils;
use utils::*;

use expect_test::expect;

#[test]
fn test_derive_empty() {
    assert_expand("DeriveEmpty", r#"struct S;"#, expect![[r#"SUBTREE $"#]]);
}

#[test]
fn test_derive_error() {
    assert_expand(
        "DeriveError",
        r#"struct S;"#,
        expect![[r##"
            SUBTREE $
              IDENT   compile_error 4294967295
              PUNCH   ! [alone] 4294967295
              SUBTREE () 4294967295
                LITERAL "#[derive(DeriveError)] struct S ;" 4294967295
              PUNCH   ; [alone] 4294967295"##]],
    );
}

#[test]
fn test_fn_like_macro_noop() {
    assert_expand(
        "fn_like_noop",
        r#"ident, 0, 1, []"#,
        expect![[r#"
            SUBTREE $
              IDENT   ident 4294967295
              PUNCH   , [alone] 4294967295
              LITERAL 0 4294967295
              PUNCH   , [alone] 4294967295
              LITERAL 1 4294967295
              PUNCH   , [alone] 4294967295
              SUBTREE [] 4294967295"#]],
    );
}

#[test]
fn test_fn_like_macro_clone_ident_subtree() {
    assert_expand(
        "fn_like_clone_tokens",
        r#"ident, []"#,
        expect![[r#"
            SUBTREE $
              IDENT   ident 4294967295
              PUNCH   , [alone] 4294967295
              SUBTREE [] 4294967295"#]],
    );
}

#[test]
fn test_fn_like_macro_clone_raw_ident() {
    assert_expand(
        "fn_like_clone_tokens",
        "r#async",
        expect![[r#"
            SUBTREE $
              IDENT   async 4294967295"#]],
    );
}

#[test]
fn test_fn_like_mk_literals() {
    assert_expand(
        "fn_like_mk_literals",
        r#""#,
        expect![[r#"
            SUBTREE $
              LITERAL b"byte_string" 4294967295
              LITERAL 'c' 4294967295
              LITERAL "string" 4294967295
              LITERAL 3.14f64 4294967295
              LITERAL 3.14 4294967295
              LITERAL 123i64 4294967295
              LITERAL 123 4294967295"#]],
    );
}

#[test]
fn test_fn_like_mk_idents() {
    // FIXME: this test is wrong: raw should be 'r#raw' but ABIs 1.64 and below
    // simply ignore `is_raw` when implementing the `Ident` interface.
    assert_expand(
        "fn_like_mk_idents",
        r#""#,
        expect![[r#"
            SUBTREE $
              IDENT   standard 4294967295
              IDENT   raw 4294967295"#]],
    );
}

#[test]
fn test_fn_like_macro_clone_literals() {
    assert_expand(
        "fn_like_clone_tokens",
        r#"1u16, 2_u32, -4i64, 3.14f32, "hello bridge""#,
        expect![[r#"
            SUBTREE $
              LITERAL 1u16 4294967295
              PUNCH   , [alone] 4294967295
              LITERAL 2_u32 4294967295
              PUNCH   , [alone] 4294967295
              PUNCH   - [alone] 4294967295
              LITERAL 4i64 4294967295
              PUNCH   , [alone] 4294967295
              LITERAL 3.14f32 4294967295
              PUNCH   , [alone] 4294967295
              LITERAL "hello bridge" 4294967295"#]],
    );
}

#[test]
fn test_attr_macro() {
    // Corresponds to
    //    #[proc_macro_test::attr_error(some arguments)]
    //    mod m {}
    assert_expand_attr(
        "attr_error",
        r#"mod m {}"#,
        r#"some arguments"#,
        expect![[r##"
            SUBTREE $
              IDENT   compile_error 4294967295
              PUNCH   ! [alone] 4294967295
              SUBTREE () 4294967295
                LITERAL "#[attr_error(some arguments)] mod m {}" 4294967295
              PUNCH   ; [alone] 4294967295"##]],
    );
}

/// Tests that we find and classify all proc macros correctly.
#[test]
fn list_test_macros() {
    let res = list().join("\n");

    expect![[r#"
        fn_like_noop [FuncLike]
        fn_like_panic [FuncLike]
        fn_like_error [FuncLike]
        fn_like_clone_tokens [FuncLike]
        fn_like_mk_literals [FuncLike]
        fn_like_mk_idents [FuncLike]
        attr_noop [Attr]
        attr_panic [Attr]
        attr_error [Attr]
        DeriveEmpty [CustomDerive]
        DerivePanic [CustomDerive]
        DeriveError [CustomDerive]"#]]
    .assert_eq(&res);
}
<|endoftext|>


<commit_before>//! Implementation of the top level `*Event` derive macro.

use proc_macro2::{Span, TokenStream};
use quote::quote;
use syn::{
    Data, DataStruct, DeriveInput, Field, Fields, FieldsNamed, Ident, Meta, MetaList, NestedMeta,
};

use crate::{
    event_parse::{to_kind_variation, EventKind, EventKindVariation},
    import_ruma_events,
};

/// Derive `Event` macro code generation.
pub fn expand_event(input: DeriveInput) -> syn::Result<TokenStream> {
    let ruma_events = import_ruma_events();

    let ident = &input.ident;
    let (kind, var) = to_kind_variation(ident).ok_or_else(|| {
        syn::Error::new(Span::call_site(), "not a valid ruma event struct identifier")
    })?;

    let fields: Vec<_> = if let Data::Struct(DataStruct { fields, .. }) = input.data.clone() {
        if let Fields::Named(FieldsNamed { named, .. }) = fields {
            if !named.iter().any(|f| f.ident.as_ref().unwrap() == "content") {
                return Err(syn::Error::new(
                    Span::call_site(),
                    "struct must contain a `content` field",
                ));
            }

            named.into_iter().collect()
        } else {
            return Err(syn::Error::new_spanned(
                fields,
                "the `Event` derive only supports named fields",
            ));
        }
    } else {
        return Err(syn::Error::new_spanned(
            input.ident,
            "the `Event` derive only supports structs with named fields",
        ));
    };

    let serialize_impl = expand_serialize_event(&input, &var, &fields, &ruma_events);
    let deserialize_impl = expand_deserialize_event(&input, &var, &fields, &ruma_events)?;
    let conversion_impl = expand_from_into(&input, &kind, &var, &fields, &ruma_events);

    let eq_impl = expand_eq_ord_event(&input, &fields);

    Ok(quote! {
        #conversion_impl
        #serialize_impl
        #deserialize_impl
        #eq_impl
    })
}

fn expand_serialize_event(
    input: &DeriveInput,
    var: &EventKindVariation,
    fields: &[Field],
    ruma_events: &TokenStream,
) -> TokenStream {
    let serde = quote! { #ruma_events::exports::serde };

    let ident = &input.ident;
    let (impl_gen, ty_gen, where_clause) = input.generics.split_for_impl();
    let serialize_fields: Vec<_> = fields
        .iter()
        .map(|field| {
            let name = field.ident.as_ref().unwrap();
            if name == "content" && var.is_redacted() {
                quote! {
                    if #ruma_events::RedactedEventContent::has_serialize_fields(&self.content) {
                        state.serialize_field("content", &self.content)?;
                    }
                }
            } else if name == "prev_content" {
                quote! {
                    if let Some(content) = self.prev_content.as_ref() {
                        state.serialize_field("prev_content", content)?;
                    }
                }
            } else if name == "unsigned" {
                quote! {
                    if !self.unsigned.is_empty() {
                        state.serialize_field("unsigned", &self.unsigned)?;
                    }
                }
            } else {
                quote! {
                    state.serialize_field(stringify!(#name), &self.#name)?;
                }
            }
        })
        .collect();

    quote! {
        #[automatically_derived]
        impl #impl_gen #serde::ser::Serialize for #ident #ty_gen #where_clause {
            fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
            where
                S: #serde::ser::Serializer,
            {
                use #serde::ser::{SerializeStruct as _, Error as _};

                let event_type = #ruma_events::EventContent::event_type(&self.content);

                let mut state = serializer.serialize_struct(stringify!(#ident), 7)?;

                state.serialize_field("type", event_type)?;
                #( #serialize_fields )*
                state.end()
            }
        }
    }
}

fn expand_deserialize_event(
    input: &DeriveInput,
    var: &EventKindVariation,
    fields: &[Field],
    ruma_events: &TokenStream,
) -> syn::Result<TokenStream> {
    let serde = quote! { #ruma_events::exports::serde };
    let serde_json = quote! { #ruma_events::exports::serde_json };

    let ident = &input.ident;
    // we know there is a content field already
    let content_type = fields
        .iter()
        // we also know that the fields are named and have an ident
        .find(|f| f.ident.as_ref().unwrap() == "content")
        .map(|f| f.ty.clone())
        .unwrap();

    let (impl_generics, ty_gen, where_clause) = input.generics.split_for_impl();
    let is_generic = !input.generics.params.is_empty();

    let enum_variants: Vec<_> = fields
        .iter()
        .map(|field| {
            let name = field.ident.as_ref().unwrap();
            to_camel_case(name)
        })
        .collect();

    let deserialize_var_types: Vec<_> = fields
        .iter()
        .map(|field| {
            let name = field.ident.as_ref().unwrap();
            let ty = &field.ty;
            if name == "content" || name == "prev_content" {
                if is_generic {
                    quote! { Box<#serde_json::value::RawValue> }
                } else {
                    quote! { #content_type }
                }
            } else {
                quote! { #ty }
            }
        })
        .collect();

    let ok_or_else_fields: Vec<_> = fields
        .iter()
        .map(|field| {
            let name = field.ident.as_ref().unwrap();
            Ok(if name == "content" {
                if is_generic && var.is_redacted() {
                    quote! {
                        let content = match C::has_deserialize_fields() {
                            #ruma_events::HasDeserializeFields::False => {
                                C::empty(&event_type).map_err(A::Error::custom)?
                            },
                            #ruma_events::HasDeserializeFields::True => {
                                let json = content.ok_or_else(
                                    || #serde::de::Error::missing_field("content"),
                                )?;
                                C::from_parts(&event_type, &json).map_err(A::Error::custom)?
                            },
                            #ruma_events::HasDeserializeFields::Optional => {
                                let json = content.unwrap_or(
                                    #serde_json::value::RawValue::from_string("{}".to_owned())
                                        .unwrap()
                                );
                                C::from_parts(&event_type, &json).map_err(A::Error::custom)?
                            },
                        };
                    }
                } else if is_generic {
                    quote! {
                        let json =
                            content.ok_or_else(|| #serde::de::Error::missing_field("content"))?;
                        let content = C::from_parts(&event_type, &json).map_err(A::Error::custom)?;
                    }
                } else {
                    quote! {
                        let content = content.ok_or_else(
                            || #serde::de::Error::missing_field("content"),
                        )?;
                    }
                }
            } else if name == "prev_content" {
                if is_generic {
                    quote! {
                        let prev_content = if let Some(json) = prev_content {
                            Some(C::from_parts(&event_type, &json).map_err(A::Error::custom)?)
                        } else {
                            None
                        };
                    }
                } else {
                    quote! {
                        let prev_content = if let Some(content) = prev_content {
                            Some(content)
                        } else {
                            None
                        };
                    }
                }
            } else if name == "unsigned" {
                quote! { let unsigned = unsigned.unwrap_or_default(); }
            } else {
                let attrs: Vec<_> = field
                    .attrs
                    .iter()
                    .filter(|a| a.path.is_ident("ruma_event"))
                    .map(|a| a.parse_meta())
                    .collect::<syn::Result<_>>()?;

                let has_default_attr = attrs.iter().any(|a| {
                    matches!(
                        a,
                        Meta::List(MetaList { nested, .. })
                        if nested.iter().any(|n| {
                            matches!(n, NestedMeta::Meta(Meta::Path(p)) if p.is_ident("default"))
                        })
                    )
                });

                if has_default_attr {
                    quote! {
                        let #name = #name.unwrap_or_default();
                    }
                } else {
                    quote! {
                        let #name = #name.ok_or_else(|| {
                            #serde::de::Error::missing_field(stringify!(#name))
                        })?;
                    }
                }
            })
        })
        .collect::<syn::Result<_>>()?;

    let field_names: Vec<_> = fields.iter().flat_map(|f| &f.ident).collect();

    let deserialize_impl_gen = if is_generic {
        let gen = &input.generics.params;
        quote! { <'de, #gen> }
    } else {
        quote! { <'de> }
    };
    let deserialize_phantom_type = if is_generic {
        quote! { ::std::marker::PhantomData }
    } else {
        quote! {}
    };

    Ok(quote! {
        #[automatically_derived]
        impl #deserialize_impl_gen #serde::de::Deserialize<'de> for #ident #ty_gen #where_clause {
            fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
            where
                D: #serde::de::Deserializer<'de>,
            {
                #[derive(#serde::Deserialize)]
                #[serde(field_identifier, rename_all = "snake_case")]
                enum Field {
                    // since this is represented as an enum we have to add it so the JSON picks it
                    // up
                    Type,
                    #( #enum_variants, )*
                    #[serde(other)]
                    Unknown,
                }

                /// Visits the fields of an event struct to handle deserialization of
                /// the `content` and `prev_content` fields.
                struct EventVisitor #impl_generics (#deserialize_phantom_type #ty_gen);

                #[automatically_derived]
                impl #deserialize_impl_gen #serde::de::Visitor<'de>
                    for EventVisitor #ty_gen #where_clause
                {
                    type Value = #ident #ty_gen;

                    fn expecting(
                        &self,
                        formatter: &mut ::std::fmt::Formatter<'_>,
                    ) -> ::std::fmt::Result {
                        write!(formatter, "struct implementing {}", stringify!(#content_type))
                    }

                    fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
                    where
                        A: #serde::de::MapAccess<'de>,
                    {
                        use #serde::de::Error as _;

                        let mut event_type: Option<String> = None;
                        #( let mut #field_names: Option<#deserialize_var_types> = None; )*

                        while let Some(key) = map.next_key()? {
                            match key {
                                Field::Unknown => {
                                    let _: #serde::de::IgnoredAny = map.next_value()?;
                                },
                                Field::Type => {
                                    if event_type.is_some() {
                                        return Err(#serde::de::Error::duplicate_field("type"));
                                    }
                                    event_type = Some(map.next_value()?);
                                }
                                #(
                                    Field::#enum_variants => {
                                        if #field_names.is_some() {
                                            return Err(#serde::de::Error::duplicate_field(
                                                stringify!(#field_names),
                                            ));
                                        }
                                        #field_names = Some(map.next_value()?);
                                    }
                                )*
                            }
                        }

                        let event_type =
                            event_type.ok_or_else(|| #serde::de::Error::missing_field("type"))?;
                        #( #ok_or_else_fields )*

                        Ok(#ident {
                            #( #field_names ),*
                        })
                    }
                }

                deserializer.deserialize_map(EventVisitor(#deserialize_phantom_type))
            }
        }
    })
}

fn expand_from_into(
    input: &DeriveInput,
    kind: &EventKind,
    var: &EventKindVariation,
    fields: &[Field],
    ruma_events: &TokenStream,
) -> Option<TokenStream> {
    let ruma_identifiers = quote! { #ruma_events::exports::ruma_identifiers };

    let ident = &input.ident;

    let (impl_generics, ty_gen, where_clause) = input.generics.split_for_impl();

    let fields: Vec<_> = fields.iter().flat_map(|f| &f.ident).collect();

    if let EventKindVariation::Sync | EventKindVariation::RedactedSync = var {
        let full_struct = kind.to_event_ident(&var.to_full_variation());
        Some(quote! {
            #[automatically_derived]
            impl #impl_generics From<#full_struct #ty_gen> for #ident #ty_gen #where_clause {
                fn from(event: #full_struct #ty_gen) -> Self {
                    let #full_struct { #( #fields, )* .. } = event;
                    Self { #( #fields, )* }
                }
            }

            #[automatically_derived]
            impl #impl_generics #ident #ty_gen #where_clause {
                /// Convert this sync event into a full event, one with a room_id field.
                pub fn into_full_event(
                    self,
                    room_id: #ruma_identifiers::RoomId,
                ) -> #full_struct #ty_gen {
                    let Self { #( #fields, )* } = self;
                    #full_struct {
                        #( #fields, )*
                        room_id,
                    }
                }
            }
        })
    } else {
        None
    }
}

fn expand_eq_ord_event(input: &DeriveInput, fields: &[Field]) -> Option<TokenStream> {
    if fields.iter().flat_map(|f| f.ident.as_ref()).any(|f| f == "event_id") {
        let ident = &input.ident;
        let (impl_gen, ty_gen, where_clause) = input.generics.split_for_impl();

        Some(quote! {
            #[automatically_derived]
            impl #impl_gen ::std::cmp::PartialEq for #ident #ty_gen #where_clause {
                /// This checks if two `EventId`s are equal.
                fn eq(&self, other: &Self) -> ::std::primitive::bool {
                    self.event_id == other.event_id
                }
            }

            #[automatically_derived]
            impl #impl_gen ::std::cmp::Eq for #ident #ty_gen #where_clause {}

            #[automatically_derived]
            impl #impl_gen ::std::cmp::PartialOrd for #ident #ty_gen #where_clause {
                /// This compares `EventId`s and orders them lexicographically.
                fn partial_cmp(&self, other: &Self) -> ::std::option::Option<::std::cmp::Ordering> {
                    self.event_id.partial_cmp(&other.event_id)
                }
            }

            #[automatically_derived]
            impl #impl_gen ::std::cmp::Ord for #ident #ty_gen #where_clause {
                /// This compares `EventId`s and orders them lexicographically.
                fn cmp(&self, other: &Self) -> ::std::cmp::Ordering {
                    self.event_id.cmp(&other.event_id)
                }
            }
        })
    } else {
        None
    }
}

/// CamelCase's a field ident like "foo_bar" to "FooBar".
fn to_camel_case(name: &Ident) -> Ident {
    let span = name.span();
    let name = name.to_string();

    let s: String = name
        .split('_')
        .map(|s| s.chars().next().unwrap().to_uppercase().to_string() + &s[1..])
        .collect();
    Ident::new(&s, span)
}
<commit_msg>events: Remove elaborate no-op<commit_after>//! Implementation of the top level `*Event` derive macro.

use proc_macro2::{Span, TokenStream};
use quote::quote;
use syn::{
    Data, DataStruct, DeriveInput, Field, Fields, FieldsNamed, Ident, Meta, MetaList, NestedMeta,
};

use crate::{
    event_parse::{to_kind_variation, EventKind, EventKindVariation},
    import_ruma_events,
};

/// Derive `Event` macro code generation.
pub fn expand_event(input: DeriveInput) -> syn::Result<TokenStream> {
    let ruma_events = import_ruma_events();

    let ident = &input.ident;
    let (kind, var) = to_kind_variation(ident).ok_or_else(|| {
        syn::Error::new(Span::call_site(), "not a valid ruma event struct identifier")
    })?;

    let fields: Vec<_> = if let Data::Struct(DataStruct { fields, .. }) = input.data.clone() {
        if let Fields::Named(FieldsNamed { named, .. }) = fields {
            if !named.iter().any(|f| f.ident.as_ref().unwrap() == "content") {
                return Err(syn::Error::new(
                    Span::call_site(),
                    "struct must contain a `content` field",
                ));
            }

            named.into_iter().collect()
        } else {
            return Err(syn::Error::new_spanned(
                fields,
                "the `Event` derive only supports named fields",
            ));
        }
    } else {
        return Err(syn::Error::new_spanned(
            input.ident,
            "the `Event` derive only supports structs with named fields",
        ));
    };

    let serialize_impl = expand_serialize_event(&input, &var, &fields, &ruma_events);
    let deserialize_impl = expand_deserialize_event(&input, &var, &fields, &ruma_events)?;
    let conversion_impl = expand_from_into(&input, &kind, &var, &fields, &ruma_events);

    let eq_impl = expand_eq_ord_event(&input, &fields);

    Ok(quote! {
        #conversion_impl
        #serialize_impl
        #deserialize_impl
        #eq_impl
    })
}

fn expand_serialize_event(
    input: &DeriveInput,
    var: &EventKindVariation,
    fields: &[Field],
    ruma_events: &TokenStream,
) -> TokenStream {
    let serde = quote! { #ruma_events::exports::serde };

    let ident = &input.ident;
    let (impl_gen, ty_gen, where_clause) = input.generics.split_for_impl();
    let serialize_fields: Vec<_> = fields
        .iter()
        .map(|field| {
            let name = field.ident.as_ref().unwrap();
            if name == "content" && var.is_redacted() {
                quote! {
                    if #ruma_events::RedactedEventContent::has_serialize_fields(&self.content) {
                        state.serialize_field("content", &self.content)?;
                    }
                }
            } else if name == "prev_content" {
                quote! {
                    if let Some(content) = self.prev_content.as_ref() {
                        state.serialize_field("prev_content", content)?;
                    }
                }
            } else if name == "unsigned" {
                quote! {
                    if !self.unsigned.is_empty() {
                        state.serialize_field("unsigned", &self.unsigned)?;
                    }
                }
            } else {
                quote! {
                    state.serialize_field(stringify!(#name), &self.#name)?;
                }
            }
        })
        .collect();

    quote! {
        #[automatically_derived]
        impl #impl_gen #serde::ser::Serialize for #ident #ty_gen #where_clause {
            fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
            where
                S: #serde::ser::Serializer,
            {
                use #serde::ser::{SerializeStruct as _, Error as _};

                let event_type = #ruma_events::EventContent::event_type(&self.content);

                let mut state = serializer.serialize_struct(stringify!(#ident), 7)?;

                state.serialize_field("type", event_type)?;
                #( #serialize_fields )*
                state.end()
            }
        }
    }
}

fn expand_deserialize_event(
    input: &DeriveInput,
    var: &EventKindVariation,
    fields: &[Field],
    ruma_events: &TokenStream,
) -> syn::Result<TokenStream> {
    let serde = quote! { #ruma_events::exports::serde };
    let serde_json = quote! { #ruma_events::exports::serde_json };

    let ident = &input.ident;
    // we know there is a content field already
    let content_type = fields
        .iter()
        // we also know that the fields are named and have an ident
        .find(|f| f.ident.as_ref().unwrap() == "content")
        .map(|f| f.ty.clone())
        .unwrap();

    let (impl_generics, ty_gen, where_clause) = input.generics.split_for_impl();
    let is_generic = !input.generics.params.is_empty();

    let enum_variants: Vec<_> = fields
        .iter()
        .map(|field| {
            let name = field.ident.as_ref().unwrap();
            to_camel_case(name)
        })
        .collect();

    let deserialize_var_types: Vec<_> = fields
        .iter()
        .map(|field| {
            let name = field.ident.as_ref().unwrap();
            let ty = &field.ty;
            if name == "content" || name == "prev_content" {
                if is_generic {
                    quote! { Box<#serde_json::value::RawValue> }
                } else {
                    quote! { #content_type }
                }
            } else {
                quote! { #ty }
            }
        })
        .collect();

    let ok_or_else_fields: Vec<_> = fields
        .iter()
        .map(|field| {
            let name = field.ident.as_ref().unwrap();
            Ok(if name == "content" {
                if is_generic && var.is_redacted() {
                    quote! {
                        let content = match C::has_deserialize_fields() {
                            #ruma_events::HasDeserializeFields::False => {
                                C::empty(&event_type).map_err(A::Error::custom)?
                            },
                            #ruma_events::HasDeserializeFields::True => {
                                let json = content.ok_or_else(
                                    || #serde::de::Error::missing_field("content"),
                                )?;
                                C::from_parts(&event_type, &json).map_err(A::Error::custom)?
                            },
                            #ruma_events::HasDeserializeFields::Optional => {
                                let json = content.unwrap_or(
                                    #serde_json::value::RawValue::from_string("{}".to_owned())
                                        .unwrap()
                                );
                                C::from_parts(&event_type, &json).map_err(A::Error::custom)?
                            },
                        };
                    }
                } else if is_generic {
                    quote! {
                        let json =
                            content.ok_or_else(|| #serde::de::Error::missing_field("content"))?;
                        let content = C::from_parts(&event_type, &json).map_err(A::Error::custom)?;
                    }
                } else {
                    quote! {
                        let content = content.ok_or_else(
                            || #serde::de::Error::missing_field("content"),
                        )?;
                    }
                }
            } else if name == "prev_content" {
                if is_generic {
                    quote! {
                        let prev_content = if let Some(json) = prev_content {
                            Some(C::from_parts(&event_type, &json).map_err(A::Error::custom)?)
                        } else {
                            None
                        };
                    }
                } else {
                    TokenStream::new()
                }
            } else if name == "unsigned" {
                quote! { let unsigned = unsigned.unwrap_or_default(); }
            } else {
                let attrs: Vec<_> = field
                    .attrs
                    .iter()
                    .filter(|a| a.path.is_ident("ruma_event"))
                    .map(|a| a.parse_meta())
                    .collect::<syn::Result<_>>()?;

                let has_default_attr = attrs.iter().any(|a| {
                    matches!(
                        a,
                        Meta::List(MetaList { nested, .. })
                        if nested.iter().any(|n| {
                            matches!(n, NestedMeta::Meta(Meta::Path(p)) if p.is_ident("default"))
                        })
                    )
                });

                if has_default_attr {
                    quote! {
                        let #name = #name.unwrap_or_default();
                    }
                } else {
                    quote! {
                        let #name = #name.ok_or_else(|| {
                            #serde::de::Error::missing_field(stringify!(#name))
                        })?;
                    }
                }
            })
        })
        .collect::<syn::Result<_>>()?;

    let field_names: Vec<_> = fields.iter().flat_map(|f| &f.ident).collect();

    let deserialize_impl_gen = if is_generic {
        let gen = &input.generics.params;
        quote! { <'de, #gen> }
    } else {
        quote! { <'de> }
    };
    let deserialize_phantom_type = if is_generic {
        quote! { ::std::marker::PhantomData }
    } else {
        quote! {}
    };

    Ok(quote! {
        #[automatically_derived]
        impl #deserialize_impl_gen #serde::de::Deserialize<'de> for #ident #ty_gen #where_clause {
            fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
            where
                D: #serde::de::Deserializer<'de>,
            {
                #[derive(#serde::Deserialize)]
                #[serde(field_identifier, rename_all = "snake_case")]
                enum Field {
                    // since this is represented as an enum we have to add it so the JSON picks it
                    // up
                    Type,
                    #( #enum_variants, )*
                    #[serde(other)]
                    Unknown,
                }

                /// Visits the fields of an event struct to handle deserialization of
                /// the `content` and `prev_content` fields.
                struct EventVisitor #impl_generics (#deserialize_phantom_type #ty_gen);

                #[automatically_derived]
                impl #deserialize_impl_gen #serde::de::Visitor<'de>
                    for EventVisitor #ty_gen #where_clause
                {
                    type Value = #ident #ty_gen;

                    fn expecting(
                        &self,
                        formatter: &mut ::std::fmt::Formatter<'_>,
                    ) -> ::std::fmt::Result {
                        write!(formatter, "struct implementing {}", stringify!(#content_type))
                    }

                    fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
                    where
                        A: #serde::de::MapAccess<'de>,
                    {
                        use #serde::de::Error as _;

                        let mut event_type: Option<String> = None;
                        #( let mut #field_names: Option<#deserialize_var_types> = None; )*

                        while let Some(key) = map.next_key()? {
                            match key {
                                Field::Unknown => {
                                    let _: #serde::de::IgnoredAny = map.next_value()?;
                                },
                                Field::Type => {
                                    if event_type.is_some() {
                                        return Err(#serde::de::Error::duplicate_field("type"));
                                    }
                                    event_type = Some(map.next_value()?);
                                }
                                #(
                                    Field::#enum_variants => {
                                        if #field_names.is_some() {
                                            return Err(#serde::de::Error::duplicate_field(
                                                stringify!(#field_names),
                                            ));
                                        }
                                        #field_names = Some(map.next_value()?);
                                    }
                                )*
                            }
                        }

                        let event_type =
                            event_type.ok_or_else(|| #serde::de::Error::missing_field("type"))?;
                        #( #ok_or_else_fields )*

                        Ok(#ident {
                            #( #field_names ),*
                        })
                    }
                }

                deserializer.deserialize_map(EventVisitor(#deserialize_phantom_type))
            }
        }
    })
}

fn expand_from_into(
    input: &DeriveInput,
    kind: &EventKind,
    var: &EventKindVariation,
    fields: &[Field],
    ruma_events: &TokenStream,
) -> Option<TokenStream> {
    let ruma_identifiers = quote! { #ruma_events::exports::ruma_identifiers };

    let ident = &input.ident;

    let (impl_generics, ty_gen, where_clause) = input.generics.split_for_impl();

    let fields: Vec<_> = fields.iter().flat_map(|f| &f.ident).collect();

    if let EventKindVariation::Sync | EventKindVariation::RedactedSync = var {
        let full_struct = kind.to_event_ident(&var.to_full_variation());
        Some(quote! {
            #[automatically_derived]
            impl #impl_generics From<#full_struct #ty_gen> for #ident #ty_gen #where_clause {
                fn from(event: #full_struct #ty_gen) -> Self {
                    let #full_struct { #( #fields, )* .. } = event;
                    Self { #( #fields, )* }
                }
            }

            #[automatically_derived]
            impl #impl_generics #ident #ty_gen #where_clause {
                /// Convert this sync event into a full event, one with a room_id field.
                pub fn into_full_event(
                    self,
                    room_id: #ruma_identifiers::RoomId,
                ) -> #full_struct #ty_gen {
                    let Self { #( #fields, )* } = self;
                    #full_struct {
                        #( #fields, )*
                        room_id,
                    }
                }
            }
        })
    } else {
        None
    }
}

fn expand_eq_ord_event(input: &DeriveInput, fields: &[Field]) -> Option<TokenStream> {
    if fields.iter().flat_map(|f| f.ident.as_ref()).any(|f| f == "event_id") {
        let ident = &input.ident;
        let (impl_gen, ty_gen, where_clause) = input.generics.split_for_impl();

        Some(quote! {
            #[automatically_derived]
            impl #impl_gen ::std::cmp::PartialEq for #ident #ty_gen #where_clause {
                /// This checks if two `EventId`s are equal.
                fn eq(&self, other: &Self) -> ::std::primitive::bool {
                    self.event_id == other.event_id
                }
            }

            #[automatically_derived]
            impl #impl_gen ::std::cmp::Eq for #ident #ty_gen #where_clause {}

            #[automatically_derived]
            impl #impl_gen ::std::cmp::PartialOrd for #ident #ty_gen #where_clause {
                /// This compares `EventId`s and orders them lexicographically.
                fn partial_cmp(&self, other: &Self) -> ::std::option::Option<::std::cmp::Ordering> {
                    self.event_id.partial_cmp(&other.event_id)
                }
            }

            #[automatically_derived]
            impl #impl_gen ::std::cmp::Ord for #ident #ty_gen #where_clause {
                /// This compares `EventId`s and orders them lexicographically.
                fn cmp(&self, other: &Self) -> ::std::cmp::Ordering {
                    self.event_id.cmp(&other.event_id)
                }
            }
        })
    } else {
        None
    }
}

/// CamelCase's a field ident like "foo_bar" to "FooBar".
fn to_camel_case(name: &Ident) -> Ident {
    let span = name.span();
    let name = name.to_string();

    let s: String = name
        .split('_')
        .map(|s| s.chars().next().unwrap().to_uppercase().to_string() + &s[1..])
        .collect();
    Ident::new(&s, span)
}
<|endoftext|>


<commit_before>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use anyhow::{bail, format_err, Result};
use aptos_config::config::DEFAULT_PORT;
use aptos_sdk::{transaction_builder::TransactionFactory, types::chain_id::ChainId};
use futures::future::join_all;
use itertools::zip;
use rand::{rngs::StdRng, Rng, SeedableRng};
use rand_core::OsRng;
use std::{
    cmp::min,
    process,
    time::{Duration, Instant},
};
use structopt::StructOpt;
use termion::color;
use transaction_emitter::{
    cluster::Cluster, instance::Instance, query_sequence_numbers, EmitJobRequest, EmitThreadParams,
    TxnEmitter,
};

#[derive(StructOpt, Debug)]
struct Args {
    #[structopt(short = "p", long, use_delimiter = true)]
    peers: Vec<String>,

    #[structopt(long, help = "If set, tries to use public peers instead of localhost")]
    vasp: bool,

    #[structopt(long)]
    health_check: bool,
    #[structopt(long)]
    emit_tx: bool,
    #[structopt(long)]
    diag: bool,

    // emit_tx options
    #[structopt(long, default_value = "15")]
    accounts_per_client: usize,
    #[structopt(long)]
    workers_per_ac: Option<usize>,
    #[structopt(long, default_value = "0")]
    wait_millis: u64,
    #[structopt(long)]
    burst: bool,
    #[structopt(long, default_value = "mint.key")]
    mint_file: String,
    #[structopt(long, default_value = "TESTING")]
    chain_id: ChainId,
    #[structopt(
        long,
        help = "Time to run --emit-tx for in seconds",
        default_value = "60"
    )]
    duration: u64,
    #[structopt(long, help = "Percentage of invalid txs", default_value = "0")]
    invalid_tx: usize,
}

#[tokio::main]
pub async fn main() {
    let args = Args::from_args();

    if !args.emit_tx && !args.diag {
        panic!("Can only use --emit-tx or --diag mode");
    }

    let util = BasicSwarmUtil::setup(&args);
    if args.diag {
        exit_on_error(util.diag(args.vasp).await);
        return;
    } else if args.emit_tx {
        exit_on_error(emit_tx(&util.cluster, &args).await);
        return;
    }

    let util = BasicSwarmUtil::setup(&args);
    exit_on_error(util.diag(args.vasp).await);
}

async fn emit_tx(cluster: &Cluster, args: &Args) -> Result<()> {
    let thread_params = EmitThreadParams {
        wait_millis: args.wait_millis,
        wait_committed: !args.burst,
    };
    let duration = Duration::from_secs(args.duration);
    let client = cluster.random_instance().rest_client();
    let mut root_account = cluster.load_aptos_root_account(&client).await?;
    let mut emitter = TxnEmitter::new(
        &mut root_account,
        client,
        TransactionFactory::new(cluster.chain_id).with_gas_unit_price(1),
        StdRng::from_seed(OsRng.gen()),
    );
    let mut emit_job_request =
        EmitJobRequest::new(cluster.all_instances().map(Instance::rest_client).collect())
            .accounts_per_client(args.accounts_per_client)
            .thread_params(thread_params)
            .invalid_transaction_ratio(args.invalid_tx);
    if let Some(workers_per_endpoint) = args.workers_per_ac {
        emit_job_request = emit_job_request.workers_per_endpoint(workers_per_endpoint);
    }
    if args.vasp {
        emit_job_request = emit_job_request.vasp();
    }
    let stats = emitter
        .emit_txn_for_with_stats(duration, emit_job_request, 10)
        .await?;
    println!("Total stats: {}", stats);
    println!("Average rate: {}", stats.rate(duration));
    Ok(())
}

fn parse_host_port(s: &str) -> Result<(String, u32, Option<u32>)> {
    let v = s.split(':').collect::<Vec<&str>>();
    if v.len() == 1 {
        let default_port = DEFAULT_PORT as u32;
        return Ok((v[0].to_string(), default_port, None));
    }
    if v.len() != 2 && v.len() != 3 {
        return Err(format_err!(
            "Failed to parse {:?} in host:port or host:port:debug_interface_port format",
            s
        ));
    }
    let host = v[0].to_string();
    let port = v[1].parse::<u32>()?;
    if v.len() == 3 {
        let debug_interface_port = v[2].parse::<u32>()?;
        return Ok((host, port, Some(debug_interface_port)));
    }
    Ok((host, port, None))
}

struct BasicSwarmUtil {
    cluster: Cluster,
}

impl BasicSwarmUtil {
    pub fn setup(args: &Args) -> Self {
        if args.peers.is_empty() {
            panic!("Peers not set in args");
        }
        let parsed_peers: Vec<_> = args
            .peers
            .iter()
            .map(|peer| parse_host_port(peer).expect("Failed to parse host_port"))
            .collect();

        let cluster =
            Cluster::from_host_port(parsed_peers, &args.mint_file, args.chain_id, args.vasp);
        Self { cluster }
    }

    pub async fn diag(&self, _vasp: bool) -> Result<()> {
        let client = self.cluster.random_instance().rest_client();
        let mut root_account = self.cluster.load_aptos_root_account(&client).await?;
        let mut faucet_account = self.cluster.load_aptos_root_account(&client).await?;
        let emitter = TxnEmitter::new(
            &mut root_account,
            client,
            TransactionFactory::new(self.cluster.chain_id).with_gas_unit_price(1),
            StdRng::from_seed(OsRng.gen()),
        );
        let faucet_account_address = faucet_account.address();
        let instances: Vec<_> = self.cluster.all_instances().collect();
        for instance in &instances {
            print!("Submitting txn through {}...", instance);
            let deadline = emitter
                .submit_single_transaction(
                    &instance.rest_client(),
                    &mut faucet_account,
                    &faucet_account_address,
                    10,
                )
                .await
                .map_err(|e| format_err!("Failed to submit txn through {}: {}", instance, e))?;
            println!("seq={}", faucet_account.sequence_number());
            println!(
                "Waiting all full nodes to get to seq {}",
                faucet_account.sequence_number()
            );
            loop {
                let addresses = &[faucet_account_address];
                let clients = instances
                    .iter()
                    .map(|instance| instance.rest_client())
                    .collect::<Vec<_>>();
                let futures = clients
                    .iter()
                    .map(|client| query_sequence_numbers(client, addresses));
                let results = join_all(futures).await;
                let mut all_good = true;
                for (instance, result) in zip(instances.iter(), results) {
                    let seq = result.map_err(|e| {
                        format_err!("Failed to query sequence number from {}: {}", instance, e)
                    })?[0];
                    let ip = instance.ip();
                    let color = if seq != faucet_account.sequence_number() {
                        all_good = false;
                        color::Fg(color::Red).to_string()
                    } else {
                        color::Fg(color::Green).to_string()
                    };
                    print!(
                        "[{}{}:{}{}]  ",
                        color,
                        &ip[..min(ip.len(), 10)],
                        seq,
                        color::Fg(color::Reset)
                    );
                }
                println!();
                if all_good {
                    break;
                }
                if Instant::now() > deadline {
                    bail!("Not all full nodes were updated and transaction expired");
                }
                tokio::time::sleep(Duration::from_secs(1)).await;
            }
        }
        println!("Looks like all full nodes are healthy!");
        Ok(())
    }
}

fn exit_on_error<T>(r: Result<T>) -> T {
    match r {
        Ok(r) => r,
        Err(err) => {
            println!("{}", err);
            process::exit(1)
        }
    }
}
<commit_msg>[TxnEmitter] add gas price<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use anyhow::{bail, format_err, Result};
use aptos_config::config::DEFAULT_PORT;
use aptos_sdk::{transaction_builder::TransactionFactory, types::chain_id::ChainId};
use futures::future::join_all;
use itertools::zip;
use rand::{rngs::StdRng, Rng, SeedableRng};
use rand_core::OsRng;
use std::{
    cmp::min,
    process,
    time::{Duration, Instant},
};
use structopt::StructOpt;
use termion::color;
use transaction_emitter::{
    cluster::Cluster, instance::Instance, query_sequence_numbers, EmitJobRequest, EmitThreadParams,
    TxnEmitter,
};

#[derive(StructOpt, Debug)]
struct Args {
    #[structopt(short = "p", long, use_delimiter = true)]
    peers: Vec<String>,

    #[structopt(long, help = "If set, tries to use public peers instead of localhost")]
    vasp: bool,

    #[structopt(long)]
    health_check: bool,
    #[structopt(long)]
    emit_tx: bool,
    #[structopt(long)]
    diag: bool,

    // emit_tx options
    #[structopt(long, default_value = "15")]
    accounts_per_client: usize,
    #[structopt(long)]
    workers_per_ac: Option<usize>,
    #[structopt(long, default_value = "0")]
    wait_millis: u64,
    #[structopt(long)]
    burst: bool,
    #[structopt(long, default_value = "mint.key")]
    mint_file: String,
    #[structopt(long, default_value = "TESTING")]
    chain_id: ChainId,
    #[structopt(
        long,
        help = "Time to run --emit-tx for in seconds",
        default_value = "60"
    )]
    duration: u64,
    #[structopt(long, help = "Percentage of invalid txs", default_value = "0")]
    invalid_tx: usize,
}

#[tokio::main]
pub async fn main() {
    let args = Args::from_args();

    if !args.emit_tx && !args.diag {
        panic!("Can only use --emit-tx or --diag mode");
    }

    let util = BasicSwarmUtil::setup(&args);
    if args.diag {
        exit_on_error(util.diag(args.vasp).await);
        return;
    } else if args.emit_tx {
        exit_on_error(emit_tx(&util.cluster, &args).await);
        return;
    }

    let util = BasicSwarmUtil::setup(&args);
    exit_on_error(util.diag(args.vasp).await);
}

async fn emit_tx(cluster: &Cluster, args: &Args) -> Result<()> {
    let thread_params = EmitThreadParams {
        wait_millis: args.wait_millis,
        wait_committed: !args.burst,
    };
    let duration = Duration::from_secs(args.duration);
    let client = cluster.random_instance().rest_client();
    let mut root_account = cluster.load_aptos_root_account(&client).await?;
    let mut emitter = TxnEmitter::new(
        &mut root_account,
        client,
        TransactionFactory::new(cluster.chain_id).with_gas_unit_price(1),
        StdRng::from_seed(OsRng.gen()),
    );
    let mut emit_job_request =
        EmitJobRequest::new(cluster.all_instances().map(Instance::rest_client).collect())
            .accounts_per_client(args.accounts_per_client)
            .thread_params(thread_params)
            .invalid_transaction_ratio(args.invalid_tx)
            .gas_price(1);
    if let Some(workers_per_endpoint) = args.workers_per_ac {
        emit_job_request = emit_job_request.workers_per_endpoint(workers_per_endpoint);
    }
    if args.vasp {
        emit_job_request = emit_job_request.vasp();
    }
    let stats = emitter
        .emit_txn_for_with_stats(duration, emit_job_request, 10)
        .await?;
    println!("Total stats: {}", stats);
    println!("Average rate: {}", stats.rate(duration));
    Ok(())
}

fn parse_host_port(s: &str) -> Result<(String, u32, Option<u32>)> {
    let v = s.split(':').collect::<Vec<&str>>();
    if v.len() == 1 {
        let default_port = DEFAULT_PORT as u32;
        return Ok((v[0].to_string(), default_port, None));
    }
    if v.len() != 2 && v.len() != 3 {
        return Err(format_err!(
            "Failed to parse {:?} in host:port or host:port:debug_interface_port format",
            s
        ));
    }
    let host = v[0].to_string();
    let port = v[1].parse::<u32>()?;
    if v.len() == 3 {
        let debug_interface_port = v[2].parse::<u32>()?;
        return Ok((host, port, Some(debug_interface_port)));
    }
    Ok((host, port, None))
}

struct BasicSwarmUtil {
    cluster: Cluster,
}

impl BasicSwarmUtil {
    pub fn setup(args: &Args) -> Self {
        if args.peers.is_empty() {
            panic!("Peers not set in args");
        }
        let parsed_peers: Vec<_> = args
            .peers
            .iter()
            .map(|peer| parse_host_port(peer).expect("Failed to parse host_port"))
            .collect();

        let cluster =
            Cluster::from_host_port(parsed_peers, &args.mint_file, args.chain_id, args.vasp);
        Self { cluster }
    }

    pub async fn diag(&self, _vasp: bool) -> Result<()> {
        let client = self.cluster.random_instance().rest_client();
        let mut root_account = self.cluster.load_aptos_root_account(&client).await?;
        let mut faucet_account = self.cluster.load_aptos_root_account(&client).await?;
        let emitter = TxnEmitter::new(
            &mut root_account,
            client,
            TransactionFactory::new(self.cluster.chain_id).with_gas_unit_price(1),
            StdRng::from_seed(OsRng.gen()),
        );
        let faucet_account_address = faucet_account.address();
        let instances: Vec<_> = self.cluster.all_instances().collect();
        for instance in &instances {
            print!("Submitting txn through {}...", instance);
            let deadline = emitter
                .submit_single_transaction(
                    &instance.rest_client(),
                    &mut faucet_account,
                    &faucet_account_address,
                    10,
                )
                .await
                .map_err(|e| format_err!("Failed to submit txn through {}: {}", instance, e))?;
            println!("seq={}", faucet_account.sequence_number());
            println!(
                "Waiting all full nodes to get to seq {}",
                faucet_account.sequence_number()
            );
            loop {
                let addresses = &[faucet_account_address];
                let clients = instances
                    .iter()
                    .map(|instance| instance.rest_client())
                    .collect::<Vec<_>>();
                let futures = clients
                    .iter()
                    .map(|client| query_sequence_numbers(client, addresses));
                let results = join_all(futures).await;
                let mut all_good = true;
                for (instance, result) in zip(instances.iter(), results) {
                    let seq = result.map_err(|e| {
                        format_err!("Failed to query sequence number from {}: {}", instance, e)
                    })?[0];
                    let ip = instance.ip();
                    let color = if seq != faucet_account.sequence_number() {
                        all_good = false;
                        color::Fg(color::Red).to_string()
                    } else {
                        color::Fg(color::Green).to_string()
                    };
                    print!(
                        "[{}{}:{}{}]  ",
                        color,
                        &ip[..min(ip.len(), 10)],
                        seq,
                        color::Fg(color::Reset)
                    );
                }
                println!();
                if all_good {
                    break;
                }
                if Instant::now() > deadline {
                    bail!("Not all full nodes were updated and transaction expired");
                }
                tokio::time::sleep(Duration::from_secs(1)).await;
            }
        }
        println!("Looks like all full nodes are healthy!");
        Ok(())
    }
}

fn exit_on_error<T>(r: Result<T>) -> T {
    match r {
        Ok(r) => r,
        Err(err) => {
            println!("{}", err);
            process::exit(1)
        }
    }
}
<|endoftext|>


<commit_before>use crate::primitive::sync::{Arc, Condvar, Mutex};
use std::fmt;

/// Enables threads to synchronize the beginning or end of some computation.
///
/// # Wait groups vs barriers
///
/// `WaitGroup` is very similar to [`Barrier`], but there are a few differences:
///
/// * [`Barrier`] needs to know the number of threads at construction, while `WaitGroup` is cloned to
///   register more threads.
///
/// * A [`Barrier`] can be reused even after all threads have synchronized, while a `WaitGroup`
///   synchronizes threads only once.
///
/// * All threads wait for others to reach the [`Barrier`]. With `WaitGroup`, each thread can choose
///   to either wait for other threads or to continue without blocking.
///
/// # Examples
///
/// ```
/// use crossbeam_utils::sync::WaitGroup;
/// use std::thread;
///
/// // Create a new wait group.
/// let wg = WaitGroup::new();
///
/// for _ in 0..4 {
///     // Create another reference to the wait group.
///     let wg = wg.clone();
///
///     thread::spawn(move || {
///         // Do some work.
///
///         // Drop the reference to the wait group.
///         drop(wg);
///     });
/// }
///
/// // Block until all threads have finished their work.
/// wg.wait();
/// # std::thread::sleep(std::time::Duration::from_millis(500)); // wait for background threads closed: https://github.com/rust-lang/miri/issues/1371
/// ```
///
/// [`Barrier`]: std::sync::Barrier
pub struct WaitGroup {
    inner: Arc<Inner>,
}

/// Inner state of a `WaitGroup`.
struct Inner {
    cvar: Condvar,
    count: Mutex<usize>,
}

impl Default for WaitGroup {
    fn default() -> Self {
        Self {
            inner: Arc::new(Inner {
                cvar: Condvar::new(),
                count: Mutex::new(1),
            }),
        }
    }
}

impl WaitGroup {
    /// Creates a new wait group and returns the single reference to it.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_utils::sync::WaitGroup;
    ///
    /// let wg = WaitGroup::new();
    /// ```
    pub fn new() -> Self {
        Self::default()
    }

    /// Drops this reference and waits until all other references are dropped.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_utils::sync::WaitGroup;
    /// use std::thread;
    ///
    /// let wg = WaitGroup::new();
    ///
    /// thread::spawn({
    ///     let wg = wg.clone();
    ///     move || {
    ///         // Block until both threads have reached `wait()`.
    ///         wg.wait();
    ///     }
    /// });
    ///
    /// // Block until both threads have reached `wait()`.
    /// wg.wait();
    /// # std::thread::sleep(std::time::Duration::from_millis(500)); // wait for background threads closed: https://github.com/rust-lang/miri/issues/1371
    /// ```
    pub fn wait(self) {
        if *self.inner.count.lock().unwrap() == 1 {
            return;
        }

        let inner = self.inner.clone();
        drop(self);

        let mut count = inner.count.lock().unwrap();
        while *count > 0 {
            count = inner.cvar.wait(count).unwrap();
        }
    }
}

impl Drop for WaitGroup {
    fn drop(&mut self) {
        let mut count = self.inner.count.lock().unwrap();
        *count -= 1;

        if *count == 0 {
            self.inner.cvar.notify_all();
        }
    }
}

impl Clone for WaitGroup {
    fn clone(&self) -> WaitGroup {
        let mut count = self.inner.count.lock().unwrap();
        *count += 1;

        WaitGroup {
            inner: self.inner.clone(),
        }
    }
}

impl fmt::Debug for WaitGroup {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let count: &usize = &*self.inner.count.lock().unwrap();
        f.debug_struct("WaitGroup").field("count", count).finish()
    }
}
<commit_msg>Fix clippy::explicit_auto_deref warning<commit_after>use crate::primitive::sync::{Arc, Condvar, Mutex};
use std::fmt;

/// Enables threads to synchronize the beginning or end of some computation.
///
/// # Wait groups vs barriers
///
/// `WaitGroup` is very similar to [`Barrier`], but there are a few differences:
///
/// * [`Barrier`] needs to know the number of threads at construction, while `WaitGroup` is cloned to
///   register more threads.
///
/// * A [`Barrier`] can be reused even after all threads have synchronized, while a `WaitGroup`
///   synchronizes threads only once.
///
/// * All threads wait for others to reach the [`Barrier`]. With `WaitGroup`, each thread can choose
///   to either wait for other threads or to continue without blocking.
///
/// # Examples
///
/// ```
/// use crossbeam_utils::sync::WaitGroup;
/// use std::thread;
///
/// // Create a new wait group.
/// let wg = WaitGroup::new();
///
/// for _ in 0..4 {
///     // Create another reference to the wait group.
///     let wg = wg.clone();
///
///     thread::spawn(move || {
///         // Do some work.
///
///         // Drop the reference to the wait group.
///         drop(wg);
///     });
/// }
///
/// // Block until all threads have finished their work.
/// wg.wait();
/// # std::thread::sleep(std::time::Duration::from_millis(500)); // wait for background threads closed: https://github.com/rust-lang/miri/issues/1371
/// ```
///
/// [`Barrier`]: std::sync::Barrier
pub struct WaitGroup {
    inner: Arc<Inner>,
}

/// Inner state of a `WaitGroup`.
struct Inner {
    cvar: Condvar,
    count: Mutex<usize>,
}

impl Default for WaitGroup {
    fn default() -> Self {
        Self {
            inner: Arc::new(Inner {
                cvar: Condvar::new(),
                count: Mutex::new(1),
            }),
        }
    }
}

impl WaitGroup {
    /// Creates a new wait group and returns the single reference to it.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_utils::sync::WaitGroup;
    ///
    /// let wg = WaitGroup::new();
    /// ```
    pub fn new() -> Self {
        Self::default()
    }

    /// Drops this reference and waits until all other references are dropped.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_utils::sync::WaitGroup;
    /// use std::thread;
    ///
    /// let wg = WaitGroup::new();
    ///
    /// thread::spawn({
    ///     let wg = wg.clone();
    ///     move || {
    ///         // Block until both threads have reached `wait()`.
    ///         wg.wait();
    ///     }
    /// });
    ///
    /// // Block until both threads have reached `wait()`.
    /// wg.wait();
    /// # std::thread::sleep(std::time::Duration::from_millis(500)); // wait for background threads closed: https://github.com/rust-lang/miri/issues/1371
    /// ```
    pub fn wait(self) {
        if *self.inner.count.lock().unwrap() == 1 {
            return;
        }

        let inner = self.inner.clone();
        drop(self);

        let mut count = inner.count.lock().unwrap();
        while *count > 0 {
            count = inner.cvar.wait(count).unwrap();
        }
    }
}

impl Drop for WaitGroup {
    fn drop(&mut self) {
        let mut count = self.inner.count.lock().unwrap();
        *count -= 1;

        if *count == 0 {
            self.inner.cvar.notify_all();
        }
    }
}

impl Clone for WaitGroup {
    fn clone(&self) -> WaitGroup {
        let mut count = self.inner.count.lock().unwrap();
        *count += 1;

        WaitGroup {
            inner: self.inner.clone(),
        }
    }
}

impl fmt::Debug for WaitGroup {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let count: &usize = &self.inner.count.lock().unwrap();
        f.debug_struct("WaitGroup").field("count", count).finish()
    }
}
<|endoftext|>