<commit_before>/**
 *   Mles server
 *
 *   Copyright (C) 2017  Juhamatti Kuusisaari / Mles developers
 *
 *   This program is free software: you can redistribute it and/or modify
 *   it under the terms of the GNU General Public License as published by
 *   the Free Software Foundation, either version 3 of the License, or
 *   (at your option) any later version.
 *
 *   This program is distributed in the hope that it will be useful,
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *   GNU General Public License for more details.
 *
 *   You should have received a copy of the GNU General Public License
 *   along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
extern crate tokio_core;
extern crate tokio_io;
extern crate futures;
extern crate mles_utils;

use std::collections::HashMap;
use std::rc::Rc;
use std::cell::RefCell;
use std::iter;
use std::io::{Error, ErrorKind};
use std::{process, env};
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::thread;

use tokio_core::net::TcpListener;
use tokio_core::net::TcpStream;
use tokio_core::reactor::Core;
use tokio_io::io;
use tokio_io::AsyncRead;

use futures::Future;
use futures::stream::{self, Stream};
use futures::sync::mpsc::unbounded;
use mles_utils::*;

const HDRL: usize = 4; //hdr len
const KEYL: usize = 8; //key len

fn main() {
    let mut peer = "".to_string();
    let mut argcnt = 0;
    for arg in env::args() {
        argcnt += 1;
        if 2 == argcnt {
            peer = arg;
            peer += ":8077";
            match peer.parse::<SocketAddr>() {
                Ok(_) => {},
                    Err(err) => {
                        println!("Error: {}\nUsage: mles [peer-address]", err);
                        process::exit(1);
                    },
            }
        }
        if argcnt > 2 {
            println!("Usage: mles [peer-address]");
            process::exit(1);
        }
    }

    let peer = match peer.parse::<SocketAddr>() {
        Ok(addr) => addr,
        Err(_) => {
            SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0)
        },
    };

    let keyval = match env::var("MLES_KEY") {
        Ok(val) => val,
        Err(_) => "".to_string(),
    };

    let mut core = Core::new().unwrap();
    let handle = core.handle();
    let address = "0.0.0.0:8077".parse().unwrap();
    let socket = match TcpListener::bind(&address, &handle) {
        Ok(listener) => listener,
        Err(err) => {
            println!("Error: {}", err);
            process::exit(1);
        },
    };
    println!("Listening on: {}", address);

    let spawned = Rc::new(RefCell::new(HashMap::new()));
    let channelmsgs = Rc::new(RefCell::new(HashMap::new()));
    let mut cnt = 0;

    let srv = socket.incoming().for_each(move |(stream, addr)| {
        println!("New Connection: {}", addr);
        let paddr = match stream.peer_addr() {
                Ok(paddr) => paddr,
                Err(_) => {
                    let addr = "0.0.0.0:0";
                    let addr = addr.parse::<SocketAddr>().unwrap();
                    addr
                }
        };
        let _val = stream.set_nodelay(true).map_err(|_| panic!("Cannot set to no delay"));;

        let (reader, writer) = stream.split();
        let (tx, rx) = unbounded();
        cnt += 1;

        let (tx_peer, rx_peer) = unbounded();

        let frame = io::read_exact(reader, vec![0;HDRL]);
        let frame = frame.and_then(move |(reader, payload)| {
            if payload.len() == 0 {
                Err(Error::new(ErrorKind::BrokenPipe, "broken pipe"))
            } else {
                if read_hdr_type(payload.as_slice()) != 'M' as u32 {
                    return Err(Error::new(ErrorKind::BrokenPipe, "incorrect header"));
                }
                let hdr_len = read_hdr_len(payload.as_slice());
                if 0 == hdr_len {
                    return Err(Error::new(ErrorKind::BrokenPipe, "incorrect header len"));
                }
                Ok((reader, hdr_len))
            }
        });

        let paddr_inner = paddr.clone();
        let keyval_inner = keyval.clone();
        let frame = frame.and_then(move |(reader, hdr_len)| {
            let tframe = io::read_exact(reader, vec![0;KEYL]);
            // verify key
            let tframe = tframe.and_then(move |(reader, key)| {
                let hkey;
                let key = read_key(key);
                if 0 == keyval_inner.len() {
                    hkey = do_hash(&paddr_inner);
                }
                else {
                    hkey = do_hash(&keyval_inner);
                }
                if hkey != key {
                    return Err(Error::new(ErrorKind::BrokenPipe, "incorrect remote key"));
                }
                Ok((reader, hdr_len))
            });
            tframe
        });

        let tx_once = tx.clone();
        let spawned_inner = spawned.clone();
        let chanmsgs_inner = channelmsgs.clone();
        let socket_once = frame.and_then(move |(reader, hdr_len)| {
            let tframe = io::read_exact(reader, vec![0;hdr_len]);
            let tframe = tframe.and_then(move |(reader, message)| {
                if 0 == message.len() {
                    return Err(Error::new(ErrorKind::BrokenPipe, "incorrect message len"));
                }
                else {
                    let mut spawned_once = spawned_inner.borrow_mut();
                    let mut chanmsgs_once = chanmsgs_inner.borrow_mut();
                    let decoded_message = message_decode(message.as_slice());
                    let channel = decoded_message.channel.clone();

                    if !spawned_once.contains_key(&channel) {
                        println!("Spawning peer channel thread");
                        thread::spawn(move || peer_conn(peer, tx_peer, tx));

                        let mut channel_entry = HashMap::new();
                        channel_entry.insert(cnt, tx_once.clone());
                        spawned_once.insert(channel.clone(), channel_entry);

                        let messages: Vec<Vec<u8>> = Vec::new();
                        chanmsgs_once.insert(channel.clone(), messages);
                    }
                    else {
                        let mut channel_entry = spawned_once.get_mut(&channel).unwrap();
                        channel_entry.insert(cnt, tx_once.clone());
                        let chanmsgs = chanmsgs_once.get(&channel).unwrap();
                        // send history to client
                        for msg in chanmsgs {
                            tx_once.send(msg.clone()).unwrap();
                        }
                    }
                    // todo: add empty message to history list too
                    println!("User {}:{} joined channel {}", cnt, decoded_message.uid, channel);
                    // forward to peer
                    Ok((reader, channel))
                }
            });
            tframe
        });

        let spawned_inner = spawned.clone();
        let chanmsgs_inner = channelmsgs.clone();
        let socket_next = socket_once.and_then(move |(reader, channel)| {
            let channel_next = channel.clone();
            let iter = stream::iter(iter::repeat(()).map(Ok::<(), Error>));
            iter.fold(reader, move |reader, _| {
                let frame = io::read_exact(reader, vec![0;HDRL]);
                let frame = frame.and_then(move |(reader, payload)| {
                    if payload.len() == 0 {
                        Err(Error::new(ErrorKind::BrokenPipe, "broken pipe"))
                    } else {
                        if read_hdr_type(payload.as_slice()) != 'M' as u32 {
                            return Err(Error::new(ErrorKind::BrokenPipe, "incorrect header"));
                        }
                        let hdr_len = read_hdr_len(payload.as_slice());
                        if 0 == hdr_len {
                            return Err(Error::new(ErrorKind::BrokenPipe, "incorrect header len"));
                        }
                        Ok((reader, payload, hdr_len))
                    }
                });

                let frame = frame.and_then(move |(reader, hdr, hdr_len)| {
                    //dummy read key
                    let tframe = io::read_exact(reader, vec![0;KEYL]);
                    let tframe = tframe.and_then(move |(reader, key)| {
                        Ok((reader, hdr, key, hdr_len))
                    });
                    tframe
                });

                let frame = frame.and_then(move |(reader, hdr, key, hdr_len)| {
                    let tframe = io::read_exact(reader, vec![0;hdr_len]);
                    let tframe = tframe.and_then(move |(reader, message)| {
                        if 0 == message.len() {
                            return Err(Error::new(ErrorKind::BrokenPipe, "incorrect message len"));
                        }
                        else {
                            Ok((reader, hdr, key, message))
                        }
                    });
                    tframe
                });

                let spawned = spawned_inner.clone();
                let chanmsgs = chanmsgs_inner.clone();
                let channel = channel_next.clone();
                frame.map(move |(reader, mut hdr, mut key, message)| {
                    key.extend(message);
                    hdr.extend(key);

                    // add to history
                    let mut channel_msgs = chanmsgs.borrow_mut();
                    let mut channel_msg = channel_msgs.get_mut(&channel).unwrap();
                    channel_msg.push(hdr.clone());

                    //distribute
                    let spawned = spawned.borrow();
                    let channels = spawned.get(&channel).unwrap();
                    for (ocnt, tx) in channels {
                        if *ocnt != cnt {
                            tx.send(hdr.clone()).unwrap();
                        }
                    }
                    reader
                })
            })
        });

        //try to get tx to peer
        let spawned_inner = spawned.clone();
        let peer_writer = rx_peer.for_each(move |(channel, peer_tx)| {
            let mut spawned_once = spawned_inner.borrow_mut();
            let mut channel_entry = spawned_once.get_mut(&channel).unwrap();
            channel_entry.insert(0, peer_tx);
            Ok(())
        });

        //let peer_writer = peer_writer.map_err(|_| ());
        let peer_writer = peer_writer.map(|_| ());

        handle.spawn(peer_writer);

        let socket_writer = rx.fold(writer, |writer, msg| {
            let amt = io::write_all(writer, msg);
            let amt = amt.map(|(writer, _)| writer);
            amt.map_err(|_| ())
        });

        //let socket_writer = socket_writer.join(peer_writer);

        let channels = spawned.clone();
        let chanmsgs = channelmsgs.clone();
        let socket_reader = socket_next.map_err(|_| ());
        let connection = socket_reader.map(|_| ()).select2(socket_writer.map(|_| ()));
        //let connection = socket_reader.map(|_| ()).select(peer_writer.map(|_| ()));
        handle.spawn(connection.then(move |_| {
            let mut chans = channels.borrow_mut();
            for (cname, channel) in chans.iter_mut() {
                if channel.contains_key(&cnt) {
                    channel.remove(&cnt);
                    if 0 == channel.len() {
                        let mut channelmsgs = chanmsgs.borrow_mut();
                        channelmsgs.remove(cname);
                        println!("Channel {} dropped.", cname);
                        drop(channel);
                    }
                    break;
                }
            }
            println!("Connection {} for user {} closed.", addr, cnt);
            Ok(())
        }));
        Ok(())
    });

    // execute server
    core.run(srv).unwrap();

    println!("Stopping server");
}

fn peer_conn(peer: SocketAddr, tx_peer_for_rcv: futures::sync::mpsc::UnboundedSender<(String, futures::sync::mpsc::UnboundedSender<Vec<u8>>)>, tx_peer: futures::sync::mpsc::UnboundedSender<Vec<u8>>) {
    let mut core = Core::new().unwrap();
    let handle = core.handle();

    let tcp = TcpStream::connect(&peer, &handle);
    //let tx_peer = Rc::new(RefCell::new(tx_peer));

    let client = tcp.and_then(move |pstream| {
        let _val = pstream.set_nodelay(true).map_err(|_| panic!("Cannot set peer to no delay"));
        println!("Successfully connected to peer");

        //save writes to db
        let (reader, writer) = pstream.split();
        let (tx, rx) = unbounded();
        match tx_peer_for_rcv.send(("Kanava".to_string(), tx)) {
            Ok(_) => {},
            Err(err) => { println!("Cannot send: {}", err); },
        };

        let frame = io::read_exact(reader, vec![0;HDRL]);
        let frame = frame.and_then(move |(reader, payload)| {
            if payload.len() == 0 {
                Err(Error::new(ErrorKind::BrokenPipe, "broken pipe"))
            } else {
                if read_hdr_type(payload.as_slice()) != 'M' as u32 {
                    return Err(Error::new(ErrorKind::BrokenPipe, "incorrect header"));
                }
                let hdr_len = read_hdr_len(payload.as_slice());
                if 0 == hdr_len {
                    return Err(Error::new(ErrorKind::BrokenPipe, "incorrect header len"));
                }
                Ok((reader, payload, hdr_len))
            }
        });

        let frame = frame.and_then(move |(reader, hdr, hdr_len)| {
            //dummy read key
            let tframe = io::read_exact(reader, vec![0;KEYL]);
            let tframe = tframe.and_then(move |(reader, key)| {
                Ok((reader, hdr, key, hdr_len))
            });
            tframe
        });

        let frame = frame.and_then(move |(reader, mut hdr, key, hdr_len)| {
            let tframe = io::read_exact(reader, vec![0;hdr_len]);
            let tframe = tframe.and_then(move |(reader, message)| {
                if 0 == message.len() {
                    return Err(Error::new(ErrorKind::BrokenPipe, "incorrect message len"));
                }
                else {
                    hdr.extend(key);
                    hdr.extend(message);
                    println!("Writing to tx");
                    tx_peer.send(hdr.clone()).unwrap();
                    Ok(())
                }
            });
            tframe
        });

        let frame = frame.map_err(|_| ());
        //let frame = frame.map(|_| ());

        //somehow get this one running
        let psocket_writer = rx.fold(writer, |writer, msg| {
            println!("Writing to rx");
            let amt = io::write_all(writer, msg);
            let amt = amt.map_err(|_| ());
            amt.map(|(writer, _)| writer)
            //amt
        });
        let psocket_writer = psocket_writer.map(|_| ());

        handle.spawn(psocket_writer);
        //let frame = frame.map_err(|_| ());

        //let frame = frame.join(psocket_writer);
        //////////////let connection = iter.map(|_| ()).select(socket_writer.map(|_| ()));
        //let connection = socket_writer.map(|_| ());
        frame.map(|_| ()).then(|_| Ok(()))
        //iter.map(|_| ()).select2(psocket_writer.map(|_| ())).then(|_| Ok(()))
        //iter.map(|_| ()).select(socket_writer.map(|_| ())).then(|_| Ok(()))
        /*handle.spawn(connection.then(|_| {
            println!("Peer connection closed");
            Ok(())
        }));
        Ok(()) */
    });

    core.run(client).unwrap();

    println!("Leaving peer channel");
}

<commit_msg>Add passing the channel.<commit_after>/**
 *   Mles server
 *
 *   Copyright (C) 2017  Juhamatti Kuusisaari / Mles developers
 *
 *   This program is free software: you can redistribute it and/or modify
 *   it under the terms of the GNU General Public License as published by
 *   the Free Software Foundation, either version 3 of the License, or
 *   (at your option) any later version.
 *
 *   This program is distributed in the hope that it will be useful,
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *   GNU General Public License for more details.
 *
 *   You should have received a copy of the GNU General Public License
 *   along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
extern crate tokio_core;
extern crate tokio_io;
extern crate futures;
extern crate mles_utils;

use std::collections::HashMap;
use std::rc::Rc;
use std::cell::RefCell;
use std::iter;
use std::io::{Error, ErrorKind};
use std::{process, env};
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::thread;

use tokio_core::net::TcpListener;
use tokio_core::net::TcpStream;
use tokio_core::reactor::Core;
use tokio_io::io;
use tokio_io::AsyncRead;

use futures::Future;
use futures::stream::{self, Stream};
use futures::sync::mpsc::unbounded;
use mles_utils::*;

const HDRL: usize = 4; //hdr len
const KEYL: usize = 8; //key len

fn main() {
    let mut peer = "".to_string();
    let mut argcnt = 0;
    for arg in env::args() {
        argcnt += 1;
        if 2 == argcnt {
            peer = arg;
            peer += ":8077";
            match peer.parse::<SocketAddr>() {
                Ok(_) => {},
                    Err(err) => {
                        println!("Error: {}\nUsage: mles [peer-address]", err);
                        process::exit(1);
                    },
            }
        }
        if argcnt > 2 {
            println!("Usage: mles [peer-address]");
            process::exit(1);
        }
    }

    let peer = match peer.parse::<SocketAddr>() {
        Ok(addr) => addr,
        Err(_) => {
            SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0)
        },
    };

    let keyval = match env::var("MLES_KEY") {
        Ok(val) => val,
        Err(_) => "".to_string(),
    };

    let mut core = Core::new().unwrap();
    let handle = core.handle();
    let address = "0.0.0.0:8077".parse().unwrap();
    let socket = match TcpListener::bind(&address, &handle) {
        Ok(listener) => listener,
        Err(err) => {
            println!("Error: {}", err);
            process::exit(1);
        },
    };
    println!("Listening on: {}", address);

    let spawned = Rc::new(RefCell::new(HashMap::new()));
    let channelmsgs = Rc::new(RefCell::new(HashMap::new()));
    let mut cnt = 0;

    let srv = socket.incoming().for_each(move |(stream, addr)| {
        println!("New Connection: {}", addr);
        let paddr = match stream.peer_addr() {
                Ok(paddr) => paddr,
                Err(_) => {
                    let addr = "0.0.0.0:0";
                    let addr = addr.parse::<SocketAddr>().unwrap();
                    addr
                }
        };
        let _val = stream.set_nodelay(true).map_err(|_| panic!("Cannot set to no delay"));;

        let (reader, writer) = stream.split();
        let (tx, rx) = unbounded();
        cnt += 1;

        let (tx_peer, rx_peer) = unbounded();

        let frame = io::read_exact(reader, vec![0;HDRL]);
        let frame = frame.and_then(move |(reader, payload)| {
            if payload.len() == 0 {
                Err(Error::new(ErrorKind::BrokenPipe, "broken pipe"))
            } else {
                if read_hdr_type(payload.as_slice()) != 'M' as u32 {
                    return Err(Error::new(ErrorKind::BrokenPipe, "incorrect header"));
                }
                let hdr_len = read_hdr_len(payload.as_slice());
                if 0 == hdr_len {
                    return Err(Error::new(ErrorKind::BrokenPipe, "incorrect header len"));
                }
                Ok((reader, hdr_len))
            }
        });

        let paddr_inner = paddr.clone();
        let keyval_inner = keyval.clone();
        let frame = frame.and_then(move |(reader, hdr_len)| {
            let tframe = io::read_exact(reader, vec![0;KEYL]);
            // verify key
            let tframe = tframe.and_then(move |(reader, key)| {
                let hkey;
                let key = read_key(key);
                if 0 == keyval_inner.len() {
                    hkey = do_hash(&paddr_inner);
                }
                else {
                    hkey = do_hash(&keyval_inner);
                }
                if hkey != key {
                    return Err(Error::new(ErrorKind::BrokenPipe, "incorrect remote key"));
                }
                Ok((reader, hdr_len))
            });
            tframe
        });

        let tx_once = tx.clone();
        let spawned_inner = spawned.clone();
        let chanmsgs_inner = channelmsgs.clone();
        let socket_once = frame.and_then(move |(reader, hdr_len)| {
            let tframe = io::read_exact(reader, vec![0;hdr_len]);
            let tframe = tframe.and_then(move |(reader, message)| {
                if 0 == message.len() {
                    return Err(Error::new(ErrorKind::BrokenPipe, "incorrect message len"));
                }
                else {
                    let mut spawned_once = spawned_inner.borrow_mut();
                    let mut chanmsgs_once = chanmsgs_inner.borrow_mut();
                    let decoded_message = message_decode(message.as_slice());
                    let channel = decoded_message.channel.clone();

                    if !spawned_once.contains_key(&channel) {
                        let chancp = channel.clone();
                        println!("Spawning peer channel thread");
                        thread::spawn(move || peer_conn(peer, chancp, tx_peer, tx));

                        let mut channel_entry = HashMap::new();
                        channel_entry.insert(cnt, tx_once.clone());
                        spawned_once.insert(channel.clone(), channel_entry);

                        let messages: Vec<Vec<u8>> = Vec::new();
                        chanmsgs_once.insert(channel.clone(), messages);
                    }
                    else {
                        let mut channel_entry = spawned_once.get_mut(&channel).unwrap();
                        channel_entry.insert(cnt, tx_once.clone());
                        let chanmsgs = chanmsgs_once.get(&channel).unwrap();
                        // send history to client
                        for msg in chanmsgs {
                            tx_once.send(msg.clone()).unwrap();
                        }
                    }
                    // todo: add empty message to history list too
                    println!("User {}:{} joined channel {}", cnt, decoded_message.uid, channel);
                    // forward to peer
                    Ok((reader, channel))
                }
            });
            tframe
        });

        let spawned_inner = spawned.clone();
        let chanmsgs_inner = channelmsgs.clone();
        let socket_next = socket_once.and_then(move |(reader, channel)| {
            let channel_next = channel.clone();
            let iter = stream::iter(iter::repeat(()).map(Ok::<(), Error>));
            iter.fold(reader, move |reader, _| {
                let frame = io::read_exact(reader, vec![0;HDRL]);
                let frame = frame.and_then(move |(reader, payload)| {
                    if payload.len() == 0 {
                        Err(Error::new(ErrorKind::BrokenPipe, "broken pipe"))
                    } else {
                        if read_hdr_type(payload.as_slice()) != 'M' as u32 {
                            return Err(Error::new(ErrorKind::BrokenPipe, "incorrect header"));
                        }
                        let hdr_len = read_hdr_len(payload.as_slice());
                        if 0 == hdr_len {
                            return Err(Error::new(ErrorKind::BrokenPipe, "incorrect header len"));
                        }
                        Ok((reader, payload, hdr_len))
                    }
                });

                let frame = frame.and_then(move |(reader, hdr, hdr_len)| {
                    //dummy read key
                    let tframe = io::read_exact(reader, vec![0;KEYL]);
                    let tframe = tframe.and_then(move |(reader, key)| {
                        Ok((reader, hdr, key, hdr_len))
                    });
                    tframe
                });

                let frame = frame.and_then(move |(reader, hdr, key, hdr_len)| {
                    let tframe = io::read_exact(reader, vec![0;hdr_len]);
                    let tframe = tframe.and_then(move |(reader, message)| {
                        if 0 == message.len() {
                            return Err(Error::new(ErrorKind::BrokenPipe, "incorrect message len"));
                        }
                        else {
                            Ok((reader, hdr, key, message))
                        }
                    });
                    tframe
                });

                let spawned = spawned_inner.clone();
                let chanmsgs = chanmsgs_inner.clone();
                let channel = channel_next.clone();
                frame.map(move |(reader, mut hdr, mut key, message)| {
                    key.extend(message);
                    hdr.extend(key);

                    // add to history
                    let mut channel_msgs = chanmsgs.borrow_mut();
                    let mut channel_msg = channel_msgs.get_mut(&channel).unwrap();
                    channel_msg.push(hdr.clone());

                    //distribute
                    let spawned = spawned.borrow();
                    let channels = spawned.get(&channel).unwrap();
                    for (ocnt, tx) in channels {
                        if *ocnt != cnt {
                            tx.send(hdr.clone()).unwrap();
                        }
                    }
                    reader
                })
            })
        });

        //try to get tx to peer
        let spawned_inner = spawned.clone();
        let peer_writer = rx_peer.for_each(move |(channel, peer_tx)| {
            let mut spawned_once = spawned_inner.borrow_mut();
            let mut channel_entry = spawned_once.get_mut(&channel).unwrap();
            channel_entry.insert(0, peer_tx);
            Ok(())
        });

        //let peer_writer = peer_writer.map_err(|_| ());
        let peer_writer = peer_writer.map(|_| ());

        handle.spawn(peer_writer);

        let socket_writer = rx.fold(writer, |writer, msg| {
            let amt = io::write_all(writer, msg);
            let amt = amt.map(|(writer, _)| writer);
            amt.map_err(|_| ())
        });

        //let socket_writer = socket_writer.join(peer_writer);

        let channels = spawned.clone();
        let chanmsgs = channelmsgs.clone();
        let socket_reader = socket_next.map_err(|_| ());
        let connection = socket_reader.map(|_| ()).select2(socket_writer.map(|_| ()));
        //let connection = socket_reader.map(|_| ()).select(peer_writer.map(|_| ()));
        handle.spawn(connection.then(move |_| {
            let mut chans = channels.borrow_mut();
            for (cname, channel) in chans.iter_mut() {
                if channel.contains_key(&cnt) {
                    channel.remove(&cnt);
                    if 0 == channel.len() {
                        let mut channelmsgs = chanmsgs.borrow_mut();
                        channelmsgs.remove(cname);
                        println!("Channel {} dropped.", cname);
                        drop(channel);
                    }
                    break;
                }
            }
            println!("Connection {} for user {} closed.", addr, cnt);
            Ok(())
        }));
        Ok(())
    });

    // execute server
    core.run(srv).unwrap();

    println!("Stopping server");
}

fn peer_conn(peer: SocketAddr, channel: String, tx_peer_for_rcv: futures::sync::mpsc::UnboundedSender<(String, futures::sync::mpsc::UnboundedSender<Vec<u8>>)>, tx_peer: futures::sync::mpsc::UnboundedSender<Vec<u8>>) {
    let mut core = Core::new().unwrap();
    let handle = core.handle();

    let tcp = TcpStream::connect(&peer, &handle);
    //let tx_peer = Rc::new(RefCell::new(tx_peer));

    let client = tcp.and_then(move |pstream| {
        let _val = pstream.set_nodelay(true).map_err(|_| panic!("Cannot set peer to no delay"));
        println!("Successfully connected to peer");

        //save writes to db
        let (reader, writer) = pstream.split();
        let (tx, rx) = unbounded();
        match tx_peer_for_rcv.send((channel, tx)) {
            Ok(_) => {},
            Err(err) => { println!("Cannot send: {}", err); },
        };

        let frame = io::read_exact(reader, vec![0;HDRL]);
        let frame = frame.and_then(move |(reader, payload)| {
            if payload.len() == 0 {
                Err(Error::new(ErrorKind::BrokenPipe, "broken pipe"))
            } else {
                if read_hdr_type(payload.as_slice()) != 'M' as u32 {
                    return Err(Error::new(ErrorKind::BrokenPipe, "incorrect header"));
                }
                let hdr_len = read_hdr_len(payload.as_slice());
                if 0 == hdr_len {
                    return Err(Error::new(ErrorKind::BrokenPipe, "incorrect header len"));
                }
                Ok((reader, payload, hdr_len))
            }
        });

        let frame = frame.and_then(move |(reader, hdr, hdr_len)| {
            //dummy read key
            let tframe = io::read_exact(reader, vec![0;KEYL]);
            let tframe = tframe.and_then(move |(reader, key)| {
                Ok((reader, hdr, key, hdr_len))
            });
            tframe
        });

        let frame = frame.and_then(move |(reader, mut hdr, key, hdr_len)| {
            let tframe = io::read_exact(reader, vec![0;hdr_len]);
            let tframe = tframe.and_then(move |(reader, message)| {
                if 0 == message.len() {
                    return Err(Error::new(ErrorKind::BrokenPipe, "incorrect message len"));
                }
                else {
                    hdr.extend(key);
                    hdr.extend(message);
                    println!("Writing to tx");
                    tx_peer.send(hdr.clone()).unwrap();
                    Ok(())
                }
            });
            tframe
        });

        let frame = frame.map_err(|_| ());
        //let frame = frame.map(|_| ());

        //somehow get this one running
        let psocket_writer = rx.fold(writer, |writer, msg| {
            println!("Writing to rx");
            let amt = io::write_all(writer, msg);
            let amt = amt.map_err(|_| ());
            amt.map(|(writer, _)| writer)
            //amt
        });
        let psocket_writer = psocket_writer.map(|_| ());

        handle.spawn(psocket_writer);
        //let frame = frame.map_err(|_| ());

        //let frame = frame.join(psocket_writer);
        //////////////let connection = iter.map(|_| ()).select(socket_writer.map(|_| ()));
        //let connection = socket_writer.map(|_| ());
        frame.map(|_| ()).then(|_| Ok(()))
        //iter.map(|_| ()).select2(psocket_writer.map(|_| ())).then(|_| Ok(()))
        //iter.map(|_| ()).select(socket_writer.map(|_| ())).then(|_| Ok(()))
        /*handle.spawn(connection.then(|_| {
            println!("Peer connection closed");
            Ok(())
        }));
        Ok(()) */
    });

    core.run(client).unwrap();

    println!("Leaving peer channel");
}

<|endoftext|>


<commit_before>// Copyright 2020 TiKV Project Authors. Licensed under Apache-2.0.

use std::collections::hash_map::Entry;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;

use futures::{Future, Sink, Stream};
use futures03::compat::Compat;
use futures03::stream::{self, StreamExt};
use grpcio::{
    DuplexSink, Error as GrpcError, RequestStream, Result as GrpcResult, RpcContext, RpcStatus,
    RpcStatusCode, WriteFlags,
};
use kvproto::cdcpb::{
    ChangeData, ChangeDataEvent, ChangeDataRequest, Compatibility, Event, ResolvedTs,
};
use protobuf::Message;
use security::{check_common_name, SecurityManager};
use tikv_util::collections::HashMap;
use tikv_util::mpsc::batch::{self, BatchReceiver, Sender as BatchSender, VecCollector};
use tikv_util::worker::*;

use crate::delegate::{Downstream, DownstreamID};
use crate::endpoint::{Deregister, Task};

static CONNECTION_ID_ALLOC: AtomicUsize = AtomicUsize::new(0);

const CDC_MSG_NOTIFY_COUNT: usize = 8;
const CDC_MAX_RESP_SIZE: usize = 6 * 1024 * 1024; // 6MB
const CDC_MSG_MAX_BATCH_SIZE: usize = 128;

/// A unique identifier of a Connection.
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub struct ConnID(usize);

impl ConnID {
    pub fn new() -> ConnID {
        ConnID(CONNECTION_ID_ALLOC.fetch_add(1, Ordering::SeqCst))
    }
}

#[derive(Debug)]
pub enum CdcEvent {
    ResolvedTs(ResolvedTs),
    Event(Event),
}

impl CdcEvent {
    pub fn size(&self) -> usize {
        match self {
            CdcEvent::ResolvedTs(r) => r.compute_size() as usize,
            CdcEvent::Event(ref e) => e.compute_size() as usize,
        }
    }
}

bitflags::bitflags! {
    pub struct FeatureGate: u8 {
        const BATCH_RESOLVED_TS = 0b00000001;
        // Uncomment when its ready.
        // const LargeTxn       = 0b00000010;
    }
}

pub struct Conn {
    id: ConnID,
    sink: BatchSender<CdcEvent>,
    downstreams: HashMap<u64, DownstreamID>,
    peer: String,
    version: Option<(semver::Version, FeatureGate)>,
}

impl Conn {
    pub fn new(sink: BatchSender<CdcEvent>, peer: String) -> Conn {
        Conn {
            id: ConnID::new(),
            sink,
            downstreams: HashMap::default(),
            version: None,
            peer,
        }
    }

    // TODO refactor into Error::Version.
    pub fn check_version_and_set_feature(&mut self, ver: semver::Version) -> Option<Compatibility> {
        // Assume batch resolved ts will be release in v4.0.7
        // For easy of testing (nightly CI), we lower the gate to v4.0.6
        // TODO bump the version when cherry pick to release branch.
        let v407_bacth_resoled_ts = semver::Version::new(4, 0, 6);

        match &self.version {
            Some((version, _)) => {
                if version == &ver {
                    None
                } else {
                    error!("different version on the same connection";
                        "previous version" => ?version, "version" => ?ver,
                        "downstream" => ?self.peer, "conn_id" => ?self.id);
                    let mut compat = Compatibility::default();
                    compat.required_version = version.to_string();
                    Some(compat)
                }
            }
            None => {
                let mut features = FeatureGate::empty();
                if v407_bacth_resoled_ts <= ver {
                    features.toggle(FeatureGate::BATCH_RESOLVED_TS);
                }
                info!("cdc connection version"; "version" => ver.to_string(), "features" => ?features);
                self.version = Some((ver, features));
                None
            }
        }
        // Return Err(Compatibility) when TiKV reaches the next major release,
        // so that we can remove feature gates.
    }

    pub fn get_feature(&self) -> Option<&FeatureGate> {
        self.version.as_ref().map(|(_, f)| f)
    }

    pub fn get_peer(&self) -> &str {
        &self.peer
    }

    pub fn get_id(&self) -> ConnID {
        self.id
    }

    pub fn take_downstreams(self) -> HashMap<u64, DownstreamID> {
        self.downstreams
    }

    pub fn get_sink(&self) -> BatchSender<CdcEvent> {
        self.sink.clone()
    }

    pub fn subscribe(&mut self, region_id: u64, downstream_id: DownstreamID) -> bool {
        match self.downstreams.entry(region_id) {
            Entry::Occupied(_) => false,
            Entry::Vacant(v) => {
                v.insert(downstream_id);
                true
            }
        }
    }

    pub fn unsubscribe(&mut self, region_id: u64) {
        self.downstreams.remove(&region_id);
    }

    pub fn downstream_id(&self, region_id: u64) -> Option<DownstreamID> {
        self.downstreams.get(&region_id).copied()
    }

    pub fn flush(&self) {
        if !self.sink.is_empty() {
            if let Some(notifier) = self.sink.get_notifier() {
                notifier.notify();
            }
        }
    }
}

/// Service implements the `ChangeData` service.
///
/// It's a front-end of the CDC service, schedules requests to the `Endpoint`.
#[derive(Clone)]
pub struct Service {
    scheduler: Scheduler<Task>,
    security_mgr: Arc<SecurityManager>,
}

impl Service {
    /// Create a ChangeData service.
    ///
    /// It requires a scheduler of an `Endpoint` in order to schedule tasks.
    pub fn new(scheduler: Scheduler<Task>, security_mgr: Arc<SecurityManager>) -> Service {
        Service {
            scheduler,
            security_mgr,
        }
    }
}

impl ChangeData for Service {
    fn event_feed(
        &mut self,
        ctx: RpcContext,
        stream: RequestStream<ChangeDataRequest>,
        sink: DuplexSink<ChangeDataEvent>,
    ) {
        if !check_common_name(self.security_mgr.cert_allowed_cn(), &ctx) {
            return;
        }
        // TODO: make it a bounded channel.
        let (tx, rx) = batch::unbounded(CDC_MSG_NOTIFY_COUNT);
        let peer = ctx.peer();
        let conn = Conn::new(tx, peer);
        let conn_id = conn.get_id();

        if let Err(status) = self
            .scheduler
            .schedule(Task::OpenConn { conn })
            .map_err(|e| RpcStatus::new(RpcStatusCode::INVALID_ARGUMENT, Some(format!("{:?}", e))))
        {
            error!("cdc connection initiate failed"; "error" => ?status);
            ctx.spawn(sink.fail(status).map_err(|e| {
                error!("cdc failed to send error"; "error" => ?e);
            }));
            return;
        }

        let peer = ctx.peer();
        let scheduler = self.scheduler.clone();
        let recv_req = stream.for_each(move |request| {
            let region_epoch = request.get_region_epoch().clone();
            let req_id = request.get_request_id();
            let version = match semver::Version::parse(request.get_header().get_ticdc_version()) {
                Ok(v) => v,
                Err(e) => {
                    warn!("empty or invalid TiCDC version, please upgrading TiCDC";
                        "version" => request.get_header().get_ticdc_version(),
                        "error" => ?e);
                    semver::Version::new(0, 0, 0)
                }
            };
            let downstream = Downstream::new(peer.clone(), region_epoch, req_id, conn_id);
            scheduler
                .schedule(Task::Register {
                    request,
                    downstream,
                    conn_id,
                    version,
                })
                .map_err(|e| {
                    GrpcError::RpcFailure(RpcStatus::new(
                        RpcStatusCode::INVALID_ARGUMENT,
                        Some(format!("{:?}", e)),
                    ))
                })
        });

        let rx = BatchReceiver::new(rx, CDC_MSG_MAX_BATCH_SIZE, Vec::new, VecCollector);
        let rx = rx
            .map(|events| {
                // The size of the response should not exceed CDC_MAX_RESP_SIZE.
                // Split the events into multiple responses by CDC_MAX_RESP_SIZE here.
                let events_len = events.len();
                let mut resp_vecs = Vec::with_capacity(events_len);
                resp_vecs.push(ChangeDataEvent::default());
                let mut current_events_size = 0;
                for event in events {
                    let event_size = event.size();
                    if current_events_size + event_size >= CDC_MAX_RESP_SIZE {
                        resp_vecs.push(ChangeDataEvent::default());
                        current_events_size = 0;
                    }
                    match event {
                        CdcEvent::Event(e) => {
                            resp_vecs.last_mut().unwrap().mut_events().push(e);
                            current_events_size += event_size;
                        }
                        CdcEvent::ResolvedTs(r) => {
                            // Set resolved ts as an individual event.
                            let events = resp_vecs.last_mut().unwrap().take_events();
                            let mut change_data_event = ChangeDataEvent::default();
                            change_data_event.set_events(events);

                            resp_vecs.last_mut().unwrap().set_resolved_ts(r);
                            resp_vecs.push(change_data_event);
                        }
                    }
                }
                let resps = resp_vecs
                    .into_iter()
                    .filter(|e| e.has_resolved_ts() || !e.events.is_empty())
                    .map(|resp| (resp, WriteFlags::default()));
                stream::iter(resps)
            })
            .flatten()
            .map(|item| GrpcResult::Ok(item));
        let send_resp = sink.send_all(Compat::new(rx));

        let peer = ctx.peer();
        let scheduler = self.scheduler.clone();
        ctx.spawn(recv_req.then(move |res| {
            // Unregister this downstream only.
            let deregister = Deregister::Conn(conn_id);
            if let Err(e) = scheduler.schedule(Task::Deregister(deregister)) {
                error!("cdc deregister failed"; "error" => ?e, "conn_id" => ?conn_id);
            }
            match res {
                Ok(_s) => {
                    info!("cdc send half closed"; "downstream" => peer, "conn_id" => ?conn_id);
                }
                Err(e) => {
                    warn!("cdc send failed"; "error" => ?e, "downstream" => peer, "conn_id" => ?conn_id);
                }
            }
            Ok(())
        }));

        let peer = ctx.peer();
        let scheduler = self.scheduler.clone();
        ctx.spawn(send_resp.then(move |res| {
            // Unregister this downstream only.
            let deregister = Deregister::Conn(conn_id);
            if let Err(e) = scheduler.schedule(Task::Deregister(deregister)) {
                error!("cdc deregister failed"; "error" => ?e);
            }
            match res {
                Ok(_s) => {
                    info!("cdc send half closed"; "downstream" => peer, "conn_id" => ?conn_id);
                }
                Err(e) => {
                    warn!("cdc send failed"; "error" => ?e, "downstream" => peer, "conn_id" => ?conn_id);
                }
            }
            Ok(())
        }));
    }
}
<commit_msg>cdc: fix the ordering of batch events (#8629)<commit_after>// Copyright 2020 TiKV Project Authors. Licensed under Apache-2.0.

use std::collections::hash_map::Entry;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;

use futures::{Future, Sink, Stream};
use futures03::compat::Compat;
use futures03::stream::{self, StreamExt};
use grpcio::{
    DuplexSink, Error as GrpcError, RequestStream, Result as GrpcResult, RpcContext, RpcStatus,
    RpcStatusCode, WriteFlags,
};
use kvproto::cdcpb::{
    ChangeData, ChangeDataEvent, ChangeDataRequest, Compatibility, Event, ResolvedTs,
};
use protobuf::Message;
use security::{check_common_name, SecurityManager};
use tikv_util::collections::HashMap;
use tikv_util::mpsc::batch::{self, BatchReceiver, Sender as BatchSender, VecCollector};
use tikv_util::worker::*;

use crate::delegate::{Downstream, DownstreamID};
use crate::endpoint::{Deregister, Task};

static CONNECTION_ID_ALLOC: AtomicUsize = AtomicUsize::new(0);

const CDC_MSG_NOTIFY_COUNT: usize = 8;
const CDC_MAX_RESP_SIZE: u32 = 6 * 1024 * 1024; // 6MB
const CDC_MSG_MAX_BATCH_SIZE: usize = 128;

/// A unique identifier of a Connection.
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub struct ConnID(usize);

impl ConnID {
    pub fn new() -> ConnID {
        ConnID(CONNECTION_ID_ALLOC.fetch_add(1, Ordering::SeqCst))
    }
}

#[derive(Clone, Debug)]
pub enum CdcEvent {
    ResolvedTs(ResolvedTs),
    Event(Event),
}

impl CdcEvent {
    pub fn size(&self) -> u32 {
        match self {
            CdcEvent::ResolvedTs(_) => 0,
            CdcEvent::Event(ref e) => e.compute_size(),
        }
    }

    pub fn event(&self) -> &Event {
        match self {
            CdcEvent::ResolvedTs(_) => unreachable!(),
            CdcEvent::Event(ref e) => e,
        }
    }

    pub fn resolved_ts(&self) -> &ResolvedTs {
        match self {
            CdcEvent::ResolvedTs(ref r) => r,
            CdcEvent::Event(_) => unreachable!(),
        }
    }
}

struct EventBatcher {
    events: Vec<ChangeDataEvent>,
    last_size: u32,
}

impl EventBatcher {
    pub fn with_capacity(cap: usize) -> EventBatcher {
        EventBatcher {
            events: Vec::with_capacity(cap),
            last_size: 0,
        }
    }

    // The size of the response should not exceed CDC_MAX_RESP_SIZE.
    // Split the events into multiple responses by CDC_MAX_RESP_SIZE here.
    pub fn push(&mut self, event: CdcEvent) {
        let size = event.size();
        if self.events.is_empty() || self.last_size + size >= CDC_MAX_RESP_SIZE {
            self.last_size = 0;
            self.events.push(ChangeDataEvent::default());
        }
        match event {
            CdcEvent::Event(e) => {
                self.last_size += size;
                self.events.last_mut().unwrap().mut_events().push(e);
            }
            CdcEvent::ResolvedTs(r) => {
                let mut change_data_event = ChangeDataEvent::default();
                change_data_event.set_resolved_ts(r);
                self.events.push(change_data_event);
                // Set last_size to MAX-1 for sending resolved ts as an individual event.
                // '-1' is to avoid empty event when the next event is still resolved ts.
                self.last_size = CDC_MAX_RESP_SIZE - 1;
            }
        }
    }

    pub fn build(self) -> Vec<ChangeDataEvent> {
        self.events
            .into_iter()
            .filter(|e| e.has_resolved_ts() || !e.events.is_empty())
            .collect()
    }
}

bitflags::bitflags! {
    pub struct FeatureGate: u8 {
        const BATCH_RESOLVED_TS = 0b00000001;
        // Uncomment when its ready.
        // const LargeTxn       = 0b00000010;
    }
}

pub struct Conn {
    id: ConnID,
    sink: BatchSender<CdcEvent>,
    downstreams: HashMap<u64, DownstreamID>,
    peer: String,
    version: Option<(semver::Version, FeatureGate)>,
}

impl Conn {
    pub fn new(sink: BatchSender<CdcEvent>, peer: String) -> Conn {
        Conn {
            id: ConnID::new(),
            sink,
            downstreams: HashMap::default(),
            version: None,
            peer,
        }
    }

    // TODO refactor into Error::Version.
    pub fn check_version_and_set_feature(&mut self, ver: semver::Version) -> Option<Compatibility> {
        // Assume batch resolved ts will be release in v4.0.7
        // For easy of testing (nightly CI), we lower the gate to v4.0.6
        // TODO bump the version when cherry pick to release branch.
        let v407_bacth_resoled_ts = semver::Version::new(4, 0, 6);

        match &self.version {
            Some((version, _)) => {
                if version == &ver {
                    None
                } else {
                    error!("different version on the same connection";
                        "previous version" => ?version, "version" => ?ver,
                        "downstream" => ?self.peer, "conn_id" => ?self.id);
                    let mut compat = Compatibility::default();
                    compat.required_version = version.to_string();
                    Some(compat)
                }
            }
            None => {
                let mut features = FeatureGate::empty();
                if v407_bacth_resoled_ts <= ver {
                    features.toggle(FeatureGate::BATCH_RESOLVED_TS);
                }
                info!("cdc connection version"; "version" => ver.to_string(), "features" => ?features);
                self.version = Some((ver, features));
                None
            }
        }
        // Return Err(Compatibility) when TiKV reaches the next major release,
        // so that we can remove feature gates.
    }

    pub fn get_feature(&self) -> Option<&FeatureGate> {
        self.version.as_ref().map(|(_, f)| f)
    }

    pub fn get_peer(&self) -> &str {
        &self.peer
    }

    pub fn get_id(&self) -> ConnID {
        self.id
    }

    pub fn take_downstreams(self) -> HashMap<u64, DownstreamID> {
        self.downstreams
    }

    pub fn get_sink(&self) -> BatchSender<CdcEvent> {
        self.sink.clone()
    }

    pub fn subscribe(&mut self, region_id: u64, downstream_id: DownstreamID) -> bool {
        match self.downstreams.entry(region_id) {
            Entry::Occupied(_) => false,
            Entry::Vacant(v) => {
                v.insert(downstream_id);
                true
            }
        }
    }

    pub fn unsubscribe(&mut self, region_id: u64) {
        self.downstreams.remove(&region_id);
    }

    pub fn downstream_id(&self, region_id: u64) -> Option<DownstreamID> {
        self.downstreams.get(&region_id).copied()
    }

    pub fn flush(&self) {
        if !self.sink.is_empty() {
            if let Some(notifier) = self.sink.get_notifier() {
                notifier.notify();
            }
        }
    }
}

/// Service implements the `ChangeData` service.
///
/// It's a front-end of the CDC service, schedules requests to the `Endpoint`.
#[derive(Clone)]
pub struct Service {
    scheduler: Scheduler<Task>,
    security_mgr: Arc<SecurityManager>,
}

impl Service {
    /// Create a ChangeData service.
    ///
    /// It requires a scheduler of an `Endpoint` in order to schedule tasks.
    pub fn new(scheduler: Scheduler<Task>, security_mgr: Arc<SecurityManager>) -> Service {
        Service {
            scheduler,
            security_mgr,
        }
    }
}

impl ChangeData for Service {
    fn event_feed(
        &mut self,
        ctx: RpcContext,
        stream: RequestStream<ChangeDataRequest>,
        sink: DuplexSink<ChangeDataEvent>,
    ) {
        if !check_common_name(self.security_mgr.cert_allowed_cn(), &ctx) {
            return;
        }
        // TODO: make it a bounded channel.
        let (tx, rx) = batch::unbounded(CDC_MSG_NOTIFY_COUNT);
        let peer = ctx.peer();
        let conn = Conn::new(tx, peer);
        let conn_id = conn.get_id();

        if let Err(status) = self
            .scheduler
            .schedule(Task::OpenConn { conn })
            .map_err(|e| RpcStatus::new(RpcStatusCode::INVALID_ARGUMENT, Some(format!("{:?}", e))))
        {
            error!("cdc connection initiate failed"; "error" => ?status);
            ctx.spawn(sink.fail(status).map_err(|e| {
                error!("cdc failed to send error"; "error" => ?e);
            }));
            return;
        }

        let peer = ctx.peer();
        let scheduler = self.scheduler.clone();
        let recv_req = stream.for_each(move |request| {
            let region_epoch = request.get_region_epoch().clone();
            let req_id = request.get_request_id();
            let version = match semver::Version::parse(request.get_header().get_ticdc_version()) {
                Ok(v) => v,
                Err(e) => {
                    warn!("empty or invalid TiCDC version, please upgrading TiCDC";
                        "version" => request.get_header().get_ticdc_version(),
                        "error" => ?e);
                    semver::Version::new(0, 0, 0)
                }
            };
            let downstream = Downstream::new(peer.clone(), region_epoch, req_id, conn_id);
            scheduler
                .schedule(Task::Register {
                    request,
                    downstream,
                    conn_id,
                    version,
                })
                .map_err(|e| {
                    GrpcError::RpcFailure(RpcStatus::new(
                        RpcStatusCode::INVALID_ARGUMENT,
                        Some(format!("{:?}", e)),
                    ))
                })
        });

        let rx = BatchReceiver::new(rx, CDC_MSG_MAX_BATCH_SIZE, Vec::new, VecCollector);
        let rx = rx
            .map(|events| {
                let mut batcher = EventBatcher::with_capacity(events.len());
                events.into_iter().for_each(|e| batcher.push(e));
                let resps = batcher
                    .build()
                    .into_iter()
                    .map(|e| (e, WriteFlags::default()));
                stream::iter(resps)
            })
            .flatten()
            .map(|item| GrpcResult::Ok(item));
        let send_resp = sink.send_all(Compat::new(rx));

        let peer = ctx.peer();
        let scheduler = self.scheduler.clone();
        ctx.spawn(recv_req.then(move |res| {
            // Unregister this downstream only.
            let deregister = Deregister::Conn(conn_id);
            if let Err(e) = scheduler.schedule(Task::Deregister(deregister)) {
                error!("cdc deregister failed"; "error" => ?e, "conn_id" => ?conn_id);
            }
            match res {
                Ok(_s) => {
                    info!("cdc send half closed"; "downstream" => peer, "conn_id" => ?conn_id);
                }
                Err(e) => {
                    warn!("cdc send failed"; "error" => ?e, "downstream" => peer, "conn_id" => ?conn_id);
                }
            }
            Ok(())
        }));

        let peer = ctx.peer();
        let scheduler = self.scheduler.clone();
        ctx.spawn(send_resp.then(move |res| {
            // Unregister this downstream only.
            let deregister = Deregister::Conn(conn_id);
            if let Err(e) = scheduler.schedule(Task::Deregister(deregister)) {
                error!("cdc deregister failed"; "error" => ?e);
            }
            match res {
                Ok(_s) => {
                    info!("cdc send half closed"; "downstream" => peer, "conn_id" => ?conn_id);
                }
                Err(e) => {
                    warn!("cdc send failed"; "error" => ?e, "downstream" => peer, "conn_id" => ?conn_id);
                }
            }
            Ok(())
        }));
    }
}

#[cfg(test)]
mod tests {
    #[cfg(feature = "prost-codec")]
    use kvproto::cdcpb::event::{
        Entries as EventEntries, Event as Event_oneof_event, Row as EventRow,
    };
    use kvproto::cdcpb::{ChangeDataEvent, Event, ResolvedTs};
    #[cfg(not(feature = "prost-codec"))]
    use kvproto::cdcpb::{EventEntries, EventRow, Event_oneof_event};

    use crate::service::{CdcEvent, EventBatcher, CDC_MAX_RESP_SIZE};

    #[test]
    fn test_event_batcher() {
        let mut batcher = EventBatcher::with_capacity(1024);

        let check_events = |result: Vec<ChangeDataEvent>, expected: Vec<Vec<CdcEvent>>| {
            assert_eq!(result.len(), expected.len());

            for i in 0..expected.len() {
                if !result[i].has_resolved_ts() {
                    assert_eq!(result[i].events.len(), expected[i].len());
                    for j in 0..expected[i].len() {
                        assert_eq!(&result[i].events[j], expected[i][j].event());
                    }
                } else {
                    assert_eq!(expected[i].len(), 1);
                    assert_eq!(result[i].get_resolved_ts(), expected[i][0].resolved_ts());
                }
            }
        };

        let mut event_small = Event::default();
        let row_small = EventRow::default();
        let mut event_entries = EventEntries::default();
        event_entries.entries = vec![row_small].into();
        event_small.event = Some(Event_oneof_event::Entries(event_entries));

        let mut event_big = Event::default();
        let mut row_big = EventRow::default();
        row_big.set_key(vec![0 as u8; CDC_MAX_RESP_SIZE as usize]);
        let mut event_entries = EventEntries::default();
        event_entries.entries = vec![row_big].into();
        event_big.event = Some(Event_oneof_event::Entries(event_entries));

        batcher.push(CdcEvent::Event(event_small.clone()));
        batcher.push(CdcEvent::ResolvedTs(ResolvedTs::default()));
        batcher.push(CdcEvent::ResolvedTs(ResolvedTs::default()));
        batcher.push(CdcEvent::Event(event_big.clone()));
        batcher.push(CdcEvent::Event(event_small.clone()));
        batcher.push(CdcEvent::Event(event_small.clone()));
        batcher.push(CdcEvent::Event(event_big.clone()));

        check_events(
            batcher.build(),
            vec![
                vec![CdcEvent::Event(event_small.clone())],
                vec![CdcEvent::ResolvedTs(ResolvedTs::default())],
                vec![CdcEvent::ResolvedTs(ResolvedTs::default())],
                vec![CdcEvent::Event(event_big.clone())],
                vec![CdcEvent::Event(event_small); 2],
                vec![CdcEvent::Event(event_big)],
            ],
        );
    }
}
<|endoftext|>


<commit_before>extern crate piston;
extern crate piston_window;
extern crate graphics;
extern crate gfx_device_gl;
extern crate glutin;
extern crate serde;
extern crate serde_hjson;
#[macro_use]
extern crate unborrow;
extern crate rodio;

mod app;
pub use app::App;

mod audio_private;
pub mod audio {
    pub use audio_private::dj::Dj;
}

mod game_object;
pub use game_object::GameObject;

#[macro_use]
pub mod component;

mod texture;
pub use texture::Texture;

mod transform;
pub use transform::Transform;

mod input_private;
pub mod input {
    pub use input_private::Axis;
    pub use input_private::controller;
    pub use input_private::input::Input;
    pub use input_private::keyboard;
    pub use input_private::mouse;
    pub use input_private::Button;
}

mod camera;
pub use camera::Camera;

mod vec2;
pub use vec2::Vec2;

pub mod physics;
<commit_msg>Add check_and_use function<commit_after>extern crate piston;
extern crate piston_window;
extern crate graphics;
extern crate gfx_device_gl;
extern crate glutin;
extern crate serde;
extern crate serde_hjson;
#[macro_use]
extern crate unborrow;
extern crate rodio;

mod app;
pub use app::App;

mod audio_private;
pub mod audio {
    pub use audio_private::dj::Dj;
}

mod game_object;
pub use game_object::GameObject;

#[macro_use]
pub mod component;

mod texture;
pub use texture::Texture;

mod transform;
pub use transform::Transform;

mod input_private;
pub mod input {
    pub use input_private::Axis;
    pub use input_private::controller;
    pub use input_private::input::Input;
    pub use input_private::keyboard;
    pub use input_private::mouse;
    pub use input_private::Button;
}

mod camera;
pub use camera::Camera;

mod vec2;
pub use vec2::Vec2;

pub mod physics;

use std::ops::SubAssign;
pub fn check_and_use<T, K>(resource: &mut T, cost: K) -> bool
where T: SubAssign<K> + PartialOrd<K>
{
    if *resource >= cost {
        *resource -= cost;
        true
    }
    else {false}
}
<|endoftext|>


<commit_before>//! `cargo upgrade`
#![warn(
    missing_docs,
    missing_debug_implementations,
    missing_copy_implementations,
    trivial_casts,
    trivial_numeric_casts,
    unsafe_code,
    unstable_features,
    unused_import_braces,
    unused_qualifications
)]

#[macro_use]
extern crate error_chain;

use crate::errors::*;
use cargo_edit::{
    find, get_latest_dependency, registry_url, update_registry_index, CrateName, Dependency,
    LocalManifest,
};
use failure::Fail;
use std::collections::{HashMap, HashSet};
use std::io::Write;
use std::path::{Path, PathBuf};
use std::process;
use structopt::StructOpt;
use termcolor::{BufferWriter, Color, ColorChoice, ColorSpec, WriteColor};
use url::Url;

mod errors {
    error_chain! {
        links {
            CargoEditLib(::cargo_edit::Error, ::cargo_edit::ErrorKind);
        }
        foreign_links {
            CargoMetadata(::failure::Compat<::cargo_metadata::Error>);
        }
    }
}

#[derive(Debug, StructOpt)]
#[structopt(bin_name = "cargo")]
enum Command {
    /// Upgrade dependencies as specified in the local manifest file (i.e. Cargo.toml).
    #[structopt(name = "upgrade", author = "")]
    #[structopt(
        after_help = "This command differs from `cargo update`, which updates the dependency versions recorded in the
local lock file (Cargo.lock).

If `<dependency>`(s) are provided, only the specified dependencies will be upgraded. The version to
upgrade to for each can be specified with e.g. `docopt@0.8.0` or `serde@>=0.9,<2.0`.

Dev, build, and all target dependencies will also be upgraded. Only dependencies from crates.io are
supported. Git/path dependencies will be ignored.

All packages in the workspace will be upgraded if the `--all` flag is supplied. The `--all` flag may
be supplied in the presence of a virtual manifest.

If the '--to-lockfile' flag is supplied, all dependencies will be upgraded to the currently locked
version as recorded in the Cargo.lock file. This flag requires that the Cargo.lock file is
up-to-date. If the lock file is missing, or it needs to be updated, cargo-upgrade will exit with an
error. If the '--to-lockfile' flag is supplied then the network won't be accessed."
    )]
    Upgrade(Args),
}

#[derive(Debug, StructOpt)]
struct Args {
    /// Crates to be upgraded.
    dependency: Vec<String>,

    /// Path to the manifest to upgrade
    #[structopt(long = "manifest-path", value_name = "path")]
    manifest_path: Option<PathBuf>,

    /// Upgrade all packages in the workspace.
    #[structopt(long = "all")]
    all: bool,

    /// Include prerelease versions when fetching from crates.io (e.g. 0.6.0-alpha').
    #[structopt(long = "allow-prerelease")]
    allow_prerelease: bool,

    /// Print changes to be made without making them.
    #[structopt(long = "dry-run")]
    dry_run: bool,

    /// Run without accessing the network
    #[structopt(long = "offline")]
    pub offline: bool,

    /// Upgrade all packages to the version in the lockfile.
    #[structopt(long = "to-lockfile", conflicts_with = "dependency")]
    pub to_lockfile: bool,
}

/// A collection of manifests.
struct Manifests(Vec<(LocalManifest, cargo_metadata::Package)>);

/// Helper function to check whether a `cargo_metadata::Dependency` is a version dependency.
fn is_version_dep(dependency: &cargo_metadata::Dependency) -> bool {
    match dependency.source {
        // This is the criterion cargo uses (in `SourceId::from_url`) to decide whether a
        // dependency has the 'registry' kind.
        Some(ref s) => s.splitn(2, '+').next() == Some("registry"),
        _ => false,
    }
}

fn dry_run_message() -> Result<()> {
    let bufwtr = BufferWriter::stdout(ColorChoice::Always);
    let mut buffer = bufwtr.buffer();
    buffer
        .set_color(ColorSpec::new().set_fg(Some(Color::Cyan)).set_bold(true))
        .chain_err(|| "Failed to set output colour")?;
    write!(&mut buffer, "Starting dry run. ").chain_err(|| "Failed to write dry run message")?;
    buffer
        .set_color(&ColorSpec::new())
        .chain_err(|| "Failed to clear output colour")?;
    writeln!(&mut buffer, "Changes will not be saved.")
        .chain_err(|| "Failed to write dry run message")?;
    bufwtr
        .print(&buffer)
        .chain_err(|| "Failed to print dry run message")
}

impl Manifests {
    /// Get all manifests in the workspace.
    fn get_all(manifest_path: &Option<PathBuf>) -> Result<Self> {
        let mut cmd = cargo_metadata::MetadataCommand::new();
        cmd.no_deps();
        if let Some(path) = manifest_path {
            cmd.manifest_path(path);
        }
        let result = cmd.exec().map_err(|e| {
            Error::from(e.compat()).chain_err(|| "Failed to get workspace metadata")
        })?;
        result
            .packages
            .into_iter()
            .map(|package| {
                Ok((
                    LocalManifest::try_new(Path::new(&package.manifest_path))?,
                    package,
                ))
            })
            .collect::<Result<Vec<_>>>()
            .map(Manifests)
    }

    /// Get the manifest specified by the manifest path. Try to make an educated guess if no path is
    /// provided.
    fn get_local_one(manifest_path: &Option<PathBuf>) -> Result<Self> {
        let resolved_manifest_path: String = find(&manifest_path)?.to_string_lossy().into();

        let manifest = LocalManifest::find(&manifest_path)?;

        let mut cmd = cargo_metadata::MetadataCommand::new();
        cmd.no_deps();
        if let Some(path) = manifest_path {
            cmd.manifest_path(path);
        }
        let result = cmd
            .exec()
            .map_err(|e| Error::from(e.compat()).chain_err(|| "Invalid manifest"))?;
        let packages = result.packages;
        let package = packages
            .iter()
            .find(|p| p.manifest_path.to_string_lossy() == resolved_manifest_path)
            // If we have successfully got metadata, but our manifest path does not correspond to a
            // package, we must have been called against a virtual manifest.
            .chain_err(|| {
                "Found virtual manifest, but this command requires running against an \
                 actual package in this workspace. Try adding `--all`."
            })?;

        Ok(Manifests(vec![(manifest, package.to_owned())]))
    }

    /// Get the the combined set of dependencies to upgrade. If the user has specified
    /// per-dependency desired versions, extract those here.
    fn get_dependencies(&self, only_update: Vec<String>) -> Result<DesiredUpgrades> {
        // Map the names of user-specified dependencies to the (optionally) requested version.
        let selected_dependencies = only_update
            .into_iter()
            .map(|name| {
                if let Some(dependency) = CrateName::new(&name.clone()).parse_as_version()? {
                    Ok((
                        dependency.name.clone(),
                        dependency.version().map(String::from),
                    ))
                } else {
                    Ok((name, None))
                }
            })
            .collect::<Result<HashMap<_, _>>>()?;

        Ok(DesiredUpgrades(
            self.0
                .iter()
                .flat_map(|&(_, ref package)| package.dependencies.clone())
                .filter(is_version_dep)
                .filter_map(|dependency| {
                    if selected_dependencies.is_empty() {
                        // User hasn't asked for any specific dependencies to be upgraded,
                        // so upgrade all the dependencies.
                        let mut dep = Dependency::new(&dependency.name);
                        if let Some(rename) = dependency.rename {
                            dep = dep.set_rename(&rename);
                        }
                        Some((dep, (dependency.registry, None)))
                    } else {
                        // User has asked for specific dependencies. Check if this dependency
                        // was specified, populating the registry from the lockfile metadata.
                        match selected_dependencies.get(&dependency.name) {
                            Some(version) => Some((
                                Dependency::new(&dependency.name),
                                (dependency.registry, version.clone()),
                            )),
                            None => None,
                        }
                    }
                })
                .collect(),
        ))
    }

    /// Upgrade the manifests on disk following the previously-determined upgrade schema.
    fn upgrade(self, upgraded_deps: &ActualUpgrades, dry_run: bool) -> Result<()> {
        if dry_run {
            dry_run_message()?;
        }

        for (mut manifest, package) in self.0 {
            println!("{}:", package.name);

            for (dep, version) in &upgraded_deps.0 {
                let mut new_dep = Dependency::new(&dep.name).set_version(version);
                if let Some(rename) = dep.rename() {
                    new_dep = new_dep.set_rename(&rename);
                }
                manifest.upgrade(&new_dep, dry_run)?;
            }
        }

        Ok(())
    }

    /// Update dependencies in Cargo.toml file(s) to match the corresponding
    /// version in Cargo.lock.
    fn sync_to_lockfile(self, dry_run: bool) -> Result<()> {
        // Get locked dependencies. For workspaces with multiple Cargo.toml
        // files, there is only a single lockfile, so it suffices to get
        // metadata for any one of Cargo.toml files.
        let (manifest, _package) =
            self.0.iter().next().ok_or_else(|| {
                ErrorKind::CargoEditLib(::cargo_edit::ErrorKind::InvalidCargoConfig)
            })?;
        let mut cmd = cargo_metadata::MetadataCommand::new();
        cmd.manifest_path(manifest.path.clone());
        cmd.other_options(vec!["--locked".to_string()]);

        let result = cmd
            .exec()
            .map_err(|e| Error::from(e.compat()).chain_err(|| "Invalid manifest"))?;

        let locked = result
            .packages
            .into_iter()
            .filter(|p| p.source.is_some()) // Source is none for local packages
            .collect::<Vec<_>>();

        if dry_run {
            dry_run_message()?;
        }

        for (mut manifest, package) in self.0 {
            println!("{}:", package.name);

            // Upgrade the manifests one at a time, as multiple manifests may
            // request the same dependency at differing versions.
            for (name, version) in package
                .dependencies
                .clone()
                .into_iter()
                .filter(is_version_dep)
                .filter_map(|d| {
                    for p in &locked {
                        // The requested dependency may be present in the lock file with different versions,
                        // but only one will be semver-compatible with the requested version.
                        if d.name == p.name && d.req.matches(&p.version) {
                            return Some((d.name, p.version.to_string()));
                        }
                    }
                    None
                })
            {
                manifest.upgrade(&Dependency::new(&name).set_version(&version), dry_run)?;
            }
        }
        Ok(())
    }
}

/// The set of dependencies to be upgraded, alongside the registries returned from cargo metadata, and
/// the desired versions, if specified by the user.
struct DesiredUpgrades(HashMap<Dependency, (Option<String>, Option<String>)>);

/// The complete specification of the upgrades that will be performed. Map of the dependency names
/// to the new versions.
struct ActualUpgrades(HashMap<Dependency, String>);

impl DesiredUpgrades {
    /// Transform the dependencies into their upgraded forms. If a version is specified, all
    /// dependencies will get that version.
    fn get_upgraded(self, allow_prerelease: bool, manifest_path: &Path) -> Result<ActualUpgrades> {
        self.0
            .into_iter()
            .map(|(dep, (registry, version))| {
                if let Some(v) = version {
                    Ok((dep, v))
                } else {
                    let registry_url = match registry {
                        Some(x) => Some(Url::parse(&x).map_err(|_| {
                            ErrorKind::CargoEditLib(::cargo_edit::ErrorKind::InvalidCargoConfig)
                        })?),
                        None => None,
                    };
                    get_latest_dependency(&dep.name, allow_prerelease, manifest_path, &registry_url)
                        .map(|new_dep| {
                            (
                                dep,
                                new_dep
                                    .version()
                                    .expect("Invalid dependency type")
                                    .to_string(),
                            )
                        })
                        .chain_err(|| "Failed to get new version")
                }
            })
            .collect::<Result<_>>()
            .map(ActualUpgrades)
    }
}

/// Main processing function. Allows us to return a `Result` so that `main` can print pretty error
/// messages.
fn process(args: Args) -> Result<()> {
    let Args {
        dependency,
        manifest_path,
        all,
        allow_prerelease,
        dry_run,
        to_lockfile,
        ..
    } = args;

    if !args.offline && !to_lockfile {
        let url = registry_url(&find(&manifest_path)?, None)?;
        update_registry_index(&url)?;
    }

    let manifests = if all {
        Manifests::get_all(&manifest_path)
    } else {
        Manifests::get_local_one(&manifest_path)
    }?;

    if to_lockfile {
        manifests.sync_to_lockfile(dry_run)
    } else {
        let existing_dependencies = manifests.get_dependencies(dependency)?;

        // Update indices for any alternative registries, unless
        // we're offline.
        if !args.offline {
            for registry_url in existing_dependencies
                .0
                .values()
                .filter_map(|(registry, _)| registry.as_ref())
                .collect::<HashSet<_>>()
            {
                update_registry_index(&Url::parse(registry_url).map_err(|_| {
                    ErrorKind::CargoEditLib(::cargo_edit::ErrorKind::InvalidCargoConfig)
                })?)?;
            }
        }

        let upgraded_dependencies =
            existing_dependencies.get_upgraded(allow_prerelease, &find(&manifest_path)?)?;

        manifests.upgrade(&upgraded_dependencies, dry_run)
    }
}

fn main() {
    let args: Command = Command::from_args();
    let Command::Upgrade(args) = args;

    if let Err(err) = process(args) {
        eprintln!("Command failed due to unhandled error: {}\n", err);

        for e in err.iter().skip(1) {
            eprintln!("Caused by: {}", e);
        }

        if let Some(backtrace) = err.backtrace() {
            eprintln!("Backtrace: {:?}", backtrace);
        }

        process::exit(1);
    }
}
<commit_msg>upgrade: infer prerelease flag from the version req<commit_after>//! `cargo upgrade`
#![warn(
    missing_docs,
    missing_debug_implementations,
    missing_copy_implementations,
    trivial_casts,
    trivial_numeric_casts,
    unsafe_code,
    unstable_features,
    unused_import_braces,
    unused_qualifications
)]

#[macro_use]
extern crate error_chain;

use crate::errors::*;
use cargo_edit::{
    find, get_latest_dependency, registry_url, update_registry_index, CrateName, Dependency,
    LocalManifest,
};
use failure::Fail;
use std::collections::{HashMap, HashSet};
use std::io::Write;
use std::path::{Path, PathBuf};
use std::process;
use structopt::StructOpt;
use termcolor::{BufferWriter, Color, ColorChoice, ColorSpec, WriteColor};
use url::Url;

mod errors {
    error_chain! {
        links {
            CargoEditLib(::cargo_edit::Error, ::cargo_edit::ErrorKind);
        }
        foreign_links {
            CargoMetadata(::failure::Compat<::cargo_metadata::Error>);
        }
    }
}

#[derive(Debug, StructOpt)]
#[structopt(bin_name = "cargo")]
enum Command {
    /// Upgrade dependencies as specified in the local manifest file (i.e. Cargo.toml).
    #[structopt(name = "upgrade", author = "")]
    #[structopt(
        after_help = "This command differs from `cargo update`, which updates the dependency versions recorded in the
local lock file (Cargo.lock).

If `<dependency>`(s) are provided, only the specified dependencies will be upgraded. The version to
upgrade to for each can be specified with e.g. `docopt@0.8.0` or `serde@>=0.9,<2.0`.

Dev, build, and all target dependencies will also be upgraded. Only dependencies from crates.io are
supported. Git/path dependencies will be ignored.

All packages in the workspace will be upgraded if the `--all` flag is supplied. The `--all` flag may
be supplied in the presence of a virtual manifest.

If the '--to-lockfile' flag is supplied, all dependencies will be upgraded to the currently locked
version as recorded in the Cargo.lock file. This flag requires that the Cargo.lock file is
up-to-date. If the lock file is missing, or it needs to be updated, cargo-upgrade will exit with an
error. If the '--to-lockfile' flag is supplied then the network won't be accessed."
    )]
    Upgrade(Args),
}

#[derive(Debug, StructOpt)]
struct Args {
    /// Crates to be upgraded.
    dependency: Vec<String>,

    /// Path to the manifest to upgrade
    #[structopt(long = "manifest-path", value_name = "path")]
    manifest_path: Option<PathBuf>,

    /// Upgrade all packages in the workspace.
    #[structopt(long = "all")]
    all: bool,

    /// Include prerelease versions when fetching from crates.io (e.g. 0.6.0-alpha').
    #[structopt(long = "allow-prerelease")]
    allow_prerelease: bool,

    /// Print changes to be made without making them.
    #[structopt(long = "dry-run")]
    dry_run: bool,

    /// Run without accessing the network
    #[structopt(long = "offline")]
    pub offline: bool,

    /// Upgrade all packages to the version in the lockfile.
    #[structopt(long = "to-lockfile", conflicts_with = "dependency")]
    pub to_lockfile: bool,
}

/// A collection of manifests.
struct Manifests(Vec<(LocalManifest, cargo_metadata::Package)>);

/// Helper function to check whether a `cargo_metadata::Dependency` is a version dependency.
fn is_version_dep(dependency: &cargo_metadata::Dependency) -> bool {
    match dependency.source {
        // This is the criterion cargo uses (in `SourceId::from_url`) to decide whether a
        // dependency has the 'registry' kind.
        Some(ref s) => s.splitn(2, '+').next() == Some("registry"),
        _ => false,
    }
}

fn dry_run_message() -> Result<()> {
    let bufwtr = BufferWriter::stdout(ColorChoice::Always);
    let mut buffer = bufwtr.buffer();
    buffer
        .set_color(ColorSpec::new().set_fg(Some(Color::Cyan)).set_bold(true))
        .chain_err(|| "Failed to set output colour")?;
    write!(&mut buffer, "Starting dry run. ").chain_err(|| "Failed to write dry run message")?;
    buffer
        .set_color(&ColorSpec::new())
        .chain_err(|| "Failed to clear output colour")?;
    writeln!(&mut buffer, "Changes will not be saved.")
        .chain_err(|| "Failed to write dry run message")?;
    bufwtr
        .print(&buffer)
        .chain_err(|| "Failed to print dry run message")
}

impl Manifests {
    /// Get all manifests in the workspace.
    fn get_all(manifest_path: &Option<PathBuf>) -> Result<Self> {
        let mut cmd = cargo_metadata::MetadataCommand::new();
        cmd.no_deps();
        if let Some(path) = manifest_path {
            cmd.manifest_path(path);
        }
        let result = cmd.exec().map_err(|e| {
            Error::from(e.compat()).chain_err(|| "Failed to get workspace metadata")
        })?;
        result
            .packages
            .into_iter()
            .map(|package| {
                Ok((
                    LocalManifest::try_new(Path::new(&package.manifest_path))?,
                    package,
                ))
            })
            .collect::<Result<Vec<_>>>()
            .map(Manifests)
    }

    /// Get the manifest specified by the manifest path. Try to make an educated guess if no path is
    /// provided.
    fn get_local_one(manifest_path: &Option<PathBuf>) -> Result<Self> {
        let resolved_manifest_path: String = find(&manifest_path)?.to_string_lossy().into();

        let manifest = LocalManifest::find(&manifest_path)?;

        let mut cmd = cargo_metadata::MetadataCommand::new();
        cmd.no_deps();
        if let Some(path) = manifest_path {
            cmd.manifest_path(path);
        }
        let result = cmd
            .exec()
            .map_err(|e| Error::from(e.compat()).chain_err(|| "Invalid manifest"))?;
        let packages = result.packages;
        let package = packages
            .iter()
            .find(|p| p.manifest_path.to_string_lossy() == resolved_manifest_path)
            // If we have successfully got metadata, but our manifest path does not correspond to a
            // package, we must have been called against a virtual manifest.
            .chain_err(|| {
                "Found virtual manifest, but this command requires running against an \
                 actual package in this workspace. Try adding `--all`."
            })?;

        Ok(Manifests(vec![(manifest, package.to_owned())]))
    }

    /// Get the the combined set of dependencies to upgrade. If the user has specified
    /// per-dependency desired versions, extract those here.
    fn get_dependencies(&self, only_update: Vec<String>) -> Result<DesiredUpgrades> {
        // Map the names of user-specified dependencies to the (optionally) requested version.
        let selected_dependencies = only_update
            .into_iter()
            .map(|name| {
                if let Some(dependency) = CrateName::new(&name.clone()).parse_as_version()? {
                    Ok((
                        dependency.name.clone(),
                        dependency.version().map(String::from),
                    ))
                } else {
                    Ok((name, None))
                }
            })
            .collect::<Result<HashMap<_, _>>>()?;

        Ok(DesiredUpgrades(
            self.0
                .iter()
                .flat_map(|&(_, ref package)| package.dependencies.clone())
                .filter(is_version_dep)
                .filter_map(|dependency| {
                    if selected_dependencies.is_empty() {
                        // User hasn't asked for any specific dependencies to be upgraded,
                        // so upgrade all the dependencies.
                        let mut dep = Dependency::new(&dependency.name);
                        if let Some(rename) = dependency.rename {
                            dep = dep.set_rename(&rename);
                        }
                        let is_prerelease = dependency.req.to_string().contains("-");
                        Some((dep, UpgradeMetadata {
                            registry: dependency.registry,
                            version: None,
                            is_prerelease
                        }))
                    } else {
                        // User has asked for specific dependencies. Check if this dependency
                        // was specified, populating the registry from the lockfile metadata.
                        match selected_dependencies.get(&dependency.name) {
                            Some(version) => Some((
                                Dependency::new(&dependency.name),
                                UpgradeMetadata {
                                    registry: dependency.registry,
                                    version: version.clone(),
                                    is_prerelease: false,
                                }
                            )),
                            None => None,
                        }
                    }
                })
                .collect(),
        ))
    }

    /// Upgrade the manifests on disk following the previously-determined upgrade schema.
    fn upgrade(self, upgraded_deps: &ActualUpgrades, dry_run: bool) -> Result<()> {
        if dry_run {
            dry_run_message()?;
        }

        for (mut manifest, package) in self.0 {
            println!("{}:", package.name);

            for (dep, version) in &upgraded_deps.0 {
                let mut new_dep = Dependency::new(&dep.name).set_version(version);
                if let Some(rename) = dep.rename() {
                    new_dep = new_dep.set_rename(&rename);
                }
                manifest.upgrade(&new_dep, dry_run)?;
            }
        }

        Ok(())
    }

    /// Update dependencies in Cargo.toml file(s) to match the corresponding
    /// version in Cargo.lock.
    fn sync_to_lockfile(self, dry_run: bool) -> Result<()> {
        // Get locked dependencies. For workspaces with multiple Cargo.toml
        // files, there is only a single lockfile, so it suffices to get
        // metadata for any one of Cargo.toml files.
        let (manifest, _package) =
            self.0.iter().next().ok_or_else(|| {
                ErrorKind::CargoEditLib(::cargo_edit::ErrorKind::InvalidCargoConfig)
            })?;
        let mut cmd = cargo_metadata::MetadataCommand::new();
        cmd.manifest_path(manifest.path.clone());
        cmd.other_options(vec!["--locked".to_string()]);

        let result = cmd
            .exec()
            .map_err(|e| Error::from(e.compat()).chain_err(|| "Invalid manifest"))?;

        let locked = result
            .packages
            .into_iter()
            .filter(|p| p.source.is_some()) // Source is none for local packages
            .collect::<Vec<_>>();

        if dry_run {
            dry_run_message()?;
        }

        for (mut manifest, package) in self.0 {
            println!("{}:", package.name);

            // Upgrade the manifests one at a time, as multiple manifests may
            // request the same dependency at differing versions.
            for (name, version) in package
                .dependencies
                .clone()
                .into_iter()
                .filter(is_version_dep)
                .filter_map(|d| {
                    for p in &locked {
                        // The requested dependency may be present in the lock file with different versions,
                        // but only one will be semver-compatible with the requested version.
                        if d.name == p.name && d.req.matches(&p.version) {
                            return Some((d.name, p.version.to_string()));
                        }
                    }
                    None
                })
            {
                manifest.upgrade(&Dependency::new(&name).set_version(&version), dry_run)?;
            }
        }
        Ok(())
    }
}

// Some metadata about the dependency
// we're trying to upgrade.
struct UpgradeMetadata {
    registry: Option<String>,
    // `Some` if the user has specified an explicit
    // version to upgrade to.
    version: Option<String>,
    is_prerelease: bool,
}

/// The set of dependencies to be upgraded, alongside the registries returned from cargo metadata, and
/// the desired versions, if specified by the user.
struct DesiredUpgrades(HashMap<Dependency, UpgradeMetadata>);

/// The complete specification of the upgrades that will be performed. Map of the dependency names
/// to the new versions.
struct ActualUpgrades(HashMap<Dependency, String>);

impl DesiredUpgrades {
    /// Transform the dependencies into their upgraded forms. If a version is specified, all
    /// dependencies will get that version.
    fn get_upgraded(self, allow_prerelease: bool, manifest_path: &Path) -> Result<ActualUpgrades> {
        self.0
            .into_iter()
            .map(|(dep, UpgradeMetadata{ registry, version, is_prerelease })| {
                if let Some(v) = version {
                    Ok((dep, v))
                } else {
                    let registry_url = match registry {
                        Some(x) => Some(Url::parse(&x).map_err(|_| {
                            ErrorKind::CargoEditLib(::cargo_edit::ErrorKind::InvalidCargoConfig)
                        })?),
                        None => None,
                    };
                    let allow_prerelease = allow_prerelease || is_prerelease;
                    get_latest_dependency(&dep.name, allow_prerelease, manifest_path, &registry_url)
                        .map(|new_dep| {
                            (
                                dep,
                                new_dep
                                    .version()
                                    .expect("Invalid dependency type")
                                    .to_string(),
                            )
                        })
                        .chain_err(|| "Failed to get new version")
                }
            })
            .collect::<Result<_>>()
            .map(ActualUpgrades)
    }
}

/// Main processing function. Allows us to return a `Result` so that `main` can print pretty error
/// messages.
fn process(args: Args) -> Result<()> {
    let Args {
        dependency,
        manifest_path,
        all,
        allow_prerelease,
        dry_run,
        to_lockfile,
        ..
    } = args;

    if !args.offline && !to_lockfile {
        let url = registry_url(&find(&manifest_path)?, None)?;
        update_registry_index(&url)?;
    }

    let manifests = if all {
        Manifests::get_all(&manifest_path)
    } else {
        Manifests::get_local_one(&manifest_path)
    }?;

    if to_lockfile {
        manifests.sync_to_lockfile(dry_run)
    } else {
        let existing_dependencies = manifests.get_dependencies(dependency)?;

        // Update indices for any alternative registries, unless
        // we're offline.
        if !args.offline {
            for registry_url in existing_dependencies
                .0
                .values()
                .filter_map(|UpgradeMetadata { registry, .. }| registry.as_ref())
                .collect::<HashSet<_>>()
            {
                update_registry_index(&Url::parse(registry_url).map_err(|_| {
                    ErrorKind::CargoEditLib(::cargo_edit::ErrorKind::InvalidCargoConfig)
                })?)?;
            }
        }

        let upgraded_dependencies =
            existing_dependencies.get_upgraded(allow_prerelease, &find(&manifest_path)?)?;

        manifests.upgrade(&upgraded_dependencies, dry_run)
    }
}

fn main() {
    let args: Command = Command::from_args();
    let Command::Upgrade(args) = args;

    if let Err(err) = process(args) {
        eprintln!("Command failed due to unhandled error: {}\n", err);

        for e in err.iter().skip(1) {
            eprintln!("Caused by: {}", e);
        }

        if let Some(backtrace) = err.backtrace() {
            eprintln!("Backtrace: {:?}", backtrace);
        }

        process::exit(1);
    }
}
<|endoftext|>