<commit_before>#![warn(rust_2018_idioms)]
#![cfg(feature = "full")]

use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::process::{Child, Command};
use tokio_test::assert_ok;

use futures::future::{self, FutureExt};
use std::env;
use std::io;
use std::process::{ExitStatus, Stdio};

fn cat() -> Command {
    let mut me = env::current_exe().unwrap();
    me.pop();

    if me.ends_with("deps") {
        me.pop();
    }

    me.push("test-cat");

    let mut cmd = Command::new(me);
    cmd.stdin(Stdio::piped()).stdout(Stdio::piped());
    cmd
}

async fn feed_cat(mut cat: Child, n: usize) -> io::Result<ExitStatus> {
    let mut stdin = cat.stdin.take().unwrap();
    let stdout = cat.stdout.take().unwrap();

    // Produce n lines on the child's stdout.
    let write = async {
        for i in 0..n {
            let bytes = format!("line {}\n", i).into_bytes();
            stdin.write_all(&bytes).await.unwrap();
        }

        drop(stdin);
    };

    let read = async {
        let mut reader = BufReader::new(stdout).lines();
        let mut num_lines = 0;

        // Try to read `n + 1` lines, ensuring the last one is empty
        // (i.e. EOF is reached after `n` lines.
        loop {
            let data = reader
                .next_line()
                .await
                .unwrap_or_else(|_| Some(String::new()))
                .expect("failed to read line");

            let num_read = data.len();
            let done = num_lines >= n;

            match (done, num_read) {
                (false, 0) => panic!("broken pipe"),
                (true, n) if n != 0 => panic!("extraneous data"),
                _ => {
                    let expected = format!("line {}", num_lines);
                    assert_eq!(expected, data);
                }
            };

            num_lines += 1;
            if num_lines >= n {
                break;
            }
        }
    };

    // Compose reading and writing concurrently.
    future::join3(write, read, cat.wait())
        .map(|(_, _, status)| status)
        .await
}

/// Check for the following properties when feeding stdin and
/// consuming stdout of a cat-like process:
///
/// - A number of lines that amounts to a number of bytes exceeding a
///   typical OS buffer size can be fed to the child without
///   deadlock. This tests that we also consume the stdout
///   concurrently; otherwise this would deadlock.
///
/// - We read the same lines from the child that we fed it.
///
/// - The child does produce EOF on stdout after the last line.
#[tokio::test]
async fn feed_a_lot() {
    let child = cat().spawn().unwrap();
    let status = feed_cat(child, 10000).await.unwrap();
    assert_eq!(status.code(), Some(0));
}

#[tokio::test]
async fn wait_with_output_captures() {
    let mut child = cat().spawn().unwrap();
    let mut stdin = child.stdin.take().unwrap();

    let write_bytes = b"1234";

    let future = async {
        stdin.write_all(write_bytes).await?;
        drop(stdin);
        let out = child.wait_with_output();
        out.await
    };

    let output = future.await.unwrap();

    assert!(output.status.success());
    assert_eq!(output.stdout, write_bytes);
    assert_eq!(output.stderr.len(), 0);
}

#[tokio::test]
async fn status_closes_any_pipes() {
    // Cat will open a pipe between the parent and child.
    // If `status_async` doesn't ensure the handles are closed,
    // we would end up blocking forever (and time out).
    let child = cat().status();

    assert_ok!(child.await);
}

#[tokio::test]
async fn try_wait() {
    let mut child = cat().spawn().unwrap();

    let id = child.id().expect("missing id");
    assert!(id > 0);

    assert_eq!(None, assert_ok!(child.try_wait()));

    // Drop the child's stdio handles so it can terminate
    drop(child.stdin.take());
    drop(child.stderr.take());
    drop(child.stdout.take());

    assert_ok!(child.wait().await);

    // test that the `.try_wait()` method is fused just like the stdlib
    assert!(assert_ok!(child.try_wait()).unwrap().success());

    // Can't get id after process has exited
    assert_eq!(child.id(), None);
}
<commit_msg>chore: use CARGO_BIN_EXE instead of inspecting env::current_exe() (#3427)<commit_after>#![warn(rust_2018_idioms)]
#![cfg(feature = "full")]

use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::process::{Child, Command};
use tokio_test::assert_ok;

use futures::future::{self, FutureExt};
use std::env;
use std::io;
use std::process::{ExitStatus, Stdio};

fn cat() -> Command {
    let mut cmd = Command::new(env!("CARGO_BIN_EXE_test-cat"));
    cmd.stdin(Stdio::piped()).stdout(Stdio::piped());
    cmd
}

async fn feed_cat(mut cat: Child, n: usize) -> io::Result<ExitStatus> {
    let mut stdin = cat.stdin.take().unwrap();
    let stdout = cat.stdout.take().unwrap();

    // Produce n lines on the child's stdout.
    let write = async {
        for i in 0..n {
            let bytes = format!("line {}\n", i).into_bytes();
            stdin.write_all(&bytes).await.unwrap();
        }

        drop(stdin);
    };

    let read = async {
        let mut reader = BufReader::new(stdout).lines();
        let mut num_lines = 0;

        // Try to read `n + 1` lines, ensuring the last one is empty
        // (i.e. EOF is reached after `n` lines.
        loop {
            let data = reader
                .next_line()
                .await
                .unwrap_or_else(|_| Some(String::new()))
                .expect("failed to read line");

            let num_read = data.len();
            let done = num_lines >= n;

            match (done, num_read) {
                (false, 0) => panic!("broken pipe"),
                (true, n) if n != 0 => panic!("extraneous data"),
                _ => {
                    let expected = format!("line {}", num_lines);
                    assert_eq!(expected, data);
                }
            };

            num_lines += 1;
            if num_lines >= n {
                break;
            }
        }
    };

    // Compose reading and writing concurrently.
    future::join3(write, read, cat.wait())
        .map(|(_, _, status)| status)
        .await
}

/// Check for the following properties when feeding stdin and
/// consuming stdout of a cat-like process:
///
/// - A number of lines that amounts to a number of bytes exceeding a
///   typical OS buffer size can be fed to the child without
///   deadlock. This tests that we also consume the stdout
///   concurrently; otherwise this would deadlock.
///
/// - We read the same lines from the child that we fed it.
///
/// - The child does produce EOF on stdout after the last line.
#[tokio::test]
async fn feed_a_lot() {
    let child = cat().spawn().unwrap();
    let status = feed_cat(child, 10000).await.unwrap();
    assert_eq!(status.code(), Some(0));
}

#[tokio::test]
async fn wait_with_output_captures() {
    let mut child = cat().spawn().unwrap();
    let mut stdin = child.stdin.take().unwrap();

    let write_bytes = b"1234";

    let future = async {
        stdin.write_all(write_bytes).await?;
        drop(stdin);
        let out = child.wait_with_output();
        out.await
    };

    let output = future.await.unwrap();

    assert!(output.status.success());
    assert_eq!(output.stdout, write_bytes);
    assert_eq!(output.stderr.len(), 0);
}

#[tokio::test]
async fn status_closes_any_pipes() {
    // Cat will open a pipe between the parent and child.
    // If `status_async` doesn't ensure the handles are closed,
    // we would end up blocking forever (and time out).
    let child = cat().status();

    assert_ok!(child.await);
}

#[tokio::test]
async fn try_wait() {
    let mut child = cat().spawn().unwrap();

    let id = child.id().expect("missing id");
    assert!(id > 0);

    assert_eq!(None, assert_ok!(child.try_wait()));

    // Drop the child's stdio handles so it can terminate
    drop(child.stdin.take());
    drop(child.stderr.take());
    drop(child.stdout.take());

    assert_ok!(child.wait().await);

    // test that the `.try_wait()` method is fused just like the stdlib
    assert!(assert_ok!(child.try_wait()).unwrap().success());

    // Can't get id after process has exited
    assert_eq!(child.id(), None);
}
<|endoftext|>


<commit_before>use codespan::{CodeMap, LineIndex, LineNumber};
use std::{fmt, io};
use termcolor::{Color, ColorSpec, WriteColor};

use crate::{Diagnostic, LabelStyle};

struct Pad<T>(T, usize);

impl<T: fmt::Display> fmt::Display for Pad<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        for _ in 0..(self.1) {
            self.0.fmt(f)?;
        }
        Ok(())
    }
}

pub fn emit<W, S>(mut writer: W, codemap: &CodeMap<S>, diagnostic: &Diagnostic) -> io::Result<()>
where
    W: WriteColor,
    S: AsRef<str>,
{
    let line_location_color = ColorSpec::new()
        // Blue is really difficult to see on the standard windows command line
        .set_fg(Some(if cfg!(windows) {
            Color::Cyan
        } else {
            Color::Blue
        }))
        .clone();
    let diagnostic_color = ColorSpec::new()
        .set_fg(Some(diagnostic.severity.color()))
        .clone();

    let highlight_color = ColorSpec::new().set_bold(true).set_intense(true).clone();

    writer.set_color(
        &highlight_color
            .clone()
            .set_fg(Some(diagnostic.severity.color())),
    )?;
    write!(writer, "{}", diagnostic.severity)?;

    if let Some(ref code) = diagnostic.code {
        write!(writer, "[{}]", code)?;
    }

    writer.set_color(&highlight_color)?;
    writeln!(writer, ": {}", diagnostic.message)?;
    writer.reset()?;

    for label in &diagnostic.labels {
        match codemap.find_file(label.span.start()) {
            None => {
                if !label.message.is_empty() {
                    writeln!(writer, "- {}", label.message)?
                }
            },
            Some(file) => {
                let (start_line, column) =
                    file.location(label.span.start()).expect("location_start");
                let (end_line, _) = file.location(label.span.end()).expect("location_end");

                writeln!(
                    writer,
                    "- {file}:{line}:{column}",
                    file = file.name(),
                    line = start_line.number(),
                    column = column.number(),
                )?;

                let start_line_span = file.line_span(start_line).expect("line_span");
                let end_line_span = file.line_span(end_line).expect("line_span");
                let line_location_width = end_line.number().to_string().len();
                let line_location_prefix = format!("{} | ", Pad(' ', line_location_width));

                let label_color = match label.style {
                    LabelStyle::Primary => diagnostic_color.clone(),
                    LabelStyle::Secondary => ColorSpec::new()
                        .set_fg(Some(if cfg!(windows) {
                            Color::Cyan
                        } else {
                            Color::Blue
                        }))
                        .clone(),
                };

                // Write prefix to marked section
                //
                writer.set_color(&line_location_color)?;
                let line_string = start_line.number().to_string();
                writeln!(writer, "{} |", Pad(' ', line_location_width))?;
                write!(writer, "{} | ", line_string)?;

                let prefix = file
                    .src_slice(start_line_span.with_end(label.span.start()))
                    .expect("prefix");
                writer.reset()?;
                write!(writer, "{}", prefix)?;

                // Write marked section
                //
                let LineNumber(start) = start_line.number();
                let LineNumber(end) = end_line.number();

                let mark_len = if start == end {
                    let marked = file.src_slice(label.span).expect("marked");
                    writer.set_color(&label_color)?;
                    write!(writer, "{}", marked)?;
                    marked.len()
                } else {
                    let marked = file
                        .src_slice(start_line_span.with_start(label.span.start()))
                        .expect("start_of_marked");
                    writer.set_color(&label_color)?;
                    write!(writer, "{}", marked)?;

                    for i in start..(end - 1) {
                        writer.set_color(&line_location_color)?;
                        write!(writer, "{} | ", (i + 1).to_string())?;

                        let line_span = file.line_span(LineIndex(i)).expect("marked_line_span");
                        let line = file
                            .src_slice(line_span)
                            .expect("line")
                            .trim_end_matches(|ch: char| ch == '\r' || ch == '\n');
                        writer.set_color(&label_color)?;
                        writeln!(writer, "{}", line)?;
                    }

                    writer.set_color(&line_location_color)?;
                    write!(writer, "{} | ", end)?;
                    let line = file
                        .src_slice(end_line_span.with_end(label.span.end()))
                        .expect("line");
                    writer.set_color(&label_color)?;
                    write!(writer, "{}", line)?;
                    line.len()
                };

                // Write suffix to marked section
                //
                let suffix = file
                    .src_slice(end_line_span.with_start(label.span.end()))
                    .expect("suffix")
                    .trim_end_matches(|ch: char| ch == '\r' || ch == '\n');
                writer.reset()?;
                writeln!(writer, "{}", suffix)?;

                // Write mark and label

                let mark = match label.style {
                    LabelStyle::Primary => '^',
                    LabelStyle::Secondary => '-',
                };

                writer.set_color(&line_location_color)?;
                write!(writer, "{}", line_location_prefix)?;
                writer.reset()?;

                writer.set_color(&label_color)?;
                write!(writer, "{}{}", Pad(' ', prefix.len()), Pad(mark, mark_len))?;
                writer.reset()?;

                if !label.message.is_empty() {
                    writer.set_color(&label_color)?;
                    writeln!(writer, " {}", label.message)?;
                    writer.reset()?;
                }
                writer.set_color(&line_location_color)?;
                writeln!(writer, "{} |", Pad(' ', line_location_width))?;
                writer.reset()?;
            },
        }
    }
    Ok(())
}
<commit_msg>Clean up blue handling<commit_after>use codespan::{CodeMap, LineIndex, LineNumber};
use std::{fmt, io};
use termcolor::{Color, ColorSpec, WriteColor};

use crate::{Diagnostic, LabelStyle};

struct Pad<T>(T, usize);

impl<T: fmt::Display> fmt::Display for Pad<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        for _ in 0..(self.1) {
            self.0.fmt(f)?;
        }
        Ok(())
    }
}

// Blue is really difficult to see on the standard windows command line
// FIXME: Make colors configurable
#[cfg(windows)]
const BLUE: Color = Color::Cyan;
#[cfg(not(windows))]
const BLUE: Color = Color::Blue;

pub fn emit<W, S>(mut writer: W, codemap: &CodeMap<S>, diagnostic: &Diagnostic) -> io::Result<()>
where
    W: WriteColor,
    S: AsRef<str>,
{
    let line_location_color = ColorSpec::new().set_fg(Some(BLUE)).clone();
    let diagnostic_color = ColorSpec::new()
        .set_fg(Some(diagnostic.severity.color()))
        .clone();

    let highlight_color = ColorSpec::new().set_bold(true).set_intense(true).clone();

    writer.set_color(
        &highlight_color
            .clone()
            .set_fg(Some(diagnostic.severity.color())),
    )?;
    write!(writer, "{}", diagnostic.severity)?;

    if let Some(ref code) = diagnostic.code {
        write!(writer, "[{}]", code)?;
    }

    writer.set_color(&highlight_color)?;
    writeln!(writer, ": {}", diagnostic.message)?;
    writer.reset()?;

    for label in &diagnostic.labels {
        match codemap.find_file(label.span.start()) {
            None => {
                if !label.message.is_empty() {
                    writeln!(writer, "- {}", label.message)?
                }
            },
            Some(file) => {
                let (start_line, column) =
                    file.location(label.span.start()).expect("location_start");
                let (end_line, _) = file.location(label.span.end()).expect("location_end");

                writeln!(
                    writer,
                    "- {file}:{line}:{column}",
                    file = file.name(),
                    line = start_line.number(),
                    column = column.number(),
                )?;

                let start_line_span = file.line_span(start_line).expect("line_span");
                let end_line_span = file.line_span(end_line).expect("line_span");
                let line_location_width = end_line.number().to_string().len();
                let line_location_prefix = format!("{} | ", Pad(' ', line_location_width));

                let label_color = match label.style {
                    LabelStyle::Primary => diagnostic_color.clone(),
                    LabelStyle::Secondary => ColorSpec::new().set_fg(Some(BLUE)).clone(),
                };

                // Write prefix to marked section
                //
                writer.set_color(&line_location_color)?;
                let line_string = start_line.number().to_string();
                writeln!(writer, "{} |", Pad(' ', line_location_width))?;
                write!(writer, "{} | ", line_string)?;

                let prefix = file
                    .src_slice(start_line_span.with_end(label.span.start()))
                    .expect("prefix");
                writer.reset()?;
                write!(writer, "{}", prefix)?;

                // Write marked section
                //
                let LineNumber(start) = start_line.number();
                let LineNumber(end) = end_line.number();

                let mark_len = if start == end {
                    let marked = file.src_slice(label.span).expect("marked");
                    writer.set_color(&label_color)?;
                    write!(writer, "{}", marked)?;
                    marked.len()
                } else {
                    let marked = file
                        .src_slice(start_line_span.with_start(label.span.start()))
                        .expect("start_of_marked");
                    writer.set_color(&label_color)?;
                    write!(writer, "{}", marked)?;

                    for i in start..(end - 1) {
                        writer.set_color(&line_location_color)?;
                        write!(writer, "{} | ", (i + 1).to_string())?;

                        let line_span = file.line_span(LineIndex(i)).expect("marked_line_span");
                        let line = file
                            .src_slice(line_span)
                            .expect("line")
                            .trim_end_matches(|ch: char| ch == '\r' || ch == '\n');
                        writer.set_color(&label_color)?;
                        writeln!(writer, "{}", line)?;
                    }

                    writer.set_color(&line_location_color)?;
                    write!(writer, "{} | ", end)?;
                    let line = file
                        .src_slice(end_line_span.with_end(label.span.end()))
                        .expect("line");
                    writer.set_color(&label_color)?;
                    write!(writer, "{}", line)?;
                    line.len()
                };

                // Write suffix to marked section
                //
                let suffix = file
                    .src_slice(end_line_span.with_start(label.span.end()))
                    .expect("suffix")
                    .trim_end_matches(|ch: char| ch == '\r' || ch == '\n');
                writer.reset()?;
                writeln!(writer, "{}", suffix)?;

                // Write mark and label

                let mark = match label.style {
                    LabelStyle::Primary => '^',
                    LabelStyle::Secondary => '-',
                };

                writer.set_color(&line_location_color)?;
                write!(writer, "{}", line_location_prefix)?;
                writer.reset()?;

                writer.set_color(&label_color)?;
                write!(writer, "{}{}", Pad(' ', prefix.len()), Pad(mark, mark_len))?;
                writer.reset()?;

                if !label.message.is_empty() {
                    writer.set_color(&label_color)?;
                    writeln!(writer, " {}", label.message)?;
                    writer.reset()?;
                }
                writer.set_color(&line_location_color)?;
                writeln!(writer, "{} |", Pad(' ', line_location_width))?;
                writer.reset()?;
            },
        }
    }
    Ok(())
}
<|endoftext|>


<commit_before>use std::time::Duration;
use std::sync::{Arc, RwLock};
use std::thread;

use tikv::raftserver::store::*;
use kvproto::raftpb::ConfChangeType;
use kvproto::metapb;
use tikv::pd::PdClient;
use tikv::util::HandyRwLock;

use super::cluster::{Cluster, Simulator};
use super::node::new_node_cluster;
use super::server::new_server_cluster;
use super::util::*;
use super::pd::TestPdClient;

fn test_simple_conf_change<T: Simulator>(cluster: &mut Cluster<T>) {
    // init_env_log();
    let r1 = cluster.bootstrap_conf_change();

    cluster.start();

    // Now region 1 only has peer (1, 1, 1);
    let (key, value) = (b"a1", b"v1");

    cluster.put(key, value);
    assert_eq!(cluster.get(key), Some(value.to_vec()));

    let engine_2 = cluster.get_engine(2);
    must_get_none(&engine_2, b"a1");
    // add peer (2,2,2) to region 1.
    cluster.change_peer(r1, ConfChangeType::AddNode, new_peer(2, 2, 2));

    let (key, value) = (b"a2", b"v2");
    cluster.put(key, value);
    assert_eq!(cluster.get(key), Some(value.to_vec()));

    // now peer 2 must have v1 and v2;
    must_get_equal(&engine_2, b"a1", b"v1");
    must_get_equal(&engine_2, b"a2", b"v2");

    let epoch = cluster.pd_client
                       .rl()
                       .get_region_by_id(cluster.id(), 1)
                       .unwrap()
                       .get_region_epoch()
                       .clone();
    let change_peer = new_admin_request(1,
                                        new_change_peer_cmd(ConfChangeType::AddNode,
                                                            new_peer(2, 2, 2),
                                                            &epoch));
    let resp = cluster.call_command_on_leader(1, change_peer, Duration::from_secs(3)).unwrap();
    assert!(resp.get_header().has_error(),
            "we can't add same peer twice");

    // Send an invalid stale epoch
    let mut stale_epoch = metapb::RegionEpoch::new();
    stale_epoch.set_version(1);
    stale_epoch.set_conf_ver(1);
    let change_peer = new_admin_request(1,
                                        new_change_peer_cmd(ConfChangeType::AddNode,
                                                            new_peer(5, 5, 5),
                                                            &stale_epoch));
    let resp = cluster.call_command_on_leader(1, change_peer, Duration::from_secs(3)).unwrap();
    assert!(resp.get_header().has_error(),
            "We can't change peer with stale epoch");

    // peer 5 must not exist
    let engine_5 = cluster.get_engine(5);
    must_get_none(&engine_5, b"a1");

    // add peer (3, 3, 3) to region 1.
    cluster.change_peer(r1, ConfChangeType::AddNode, new_peer(3, 3, 3));
    // Remove peer (2, 2, 2) from region 1.
    cluster.change_peer(r1, ConfChangeType::RemoveNode, new_peer(2, 2, 2));

    let (key, value) = (b"a3", b"v3");
    cluster.put(key, value);
    assert_eq!(cluster.get(key), Some(value.to_vec()));
    // now peer 3 must have v1, v2 and v3
    let engine_3 = cluster.get_engine(3);
    must_get_equal(&engine_3, b"a1", b"v1");
    must_get_equal(&engine_3, b"a2", b"v2");
    must_get_equal(&engine_3, b"a3", b"v3");

    // peer 2 has nothing
    must_get_none(&engine_2, b"a1");
    must_get_none(&engine_2, b"a2");

    let epoch = cluster.pd_client
                       .rl()
                       .get_region_by_id(cluster.id(), 1)
                       .unwrap()
                       .get_region_epoch()
                       .clone();
    let change_peer = new_admin_request(1,
                                        new_change_peer_cmd(ConfChangeType::RemoveNode,
                                                            new_peer(2, 2, 2),
                                                            &epoch));
    let resp = cluster.call_command_on_leader(1, change_peer, Duration::from_secs(3)).unwrap();
    assert!(resp.get_header().has_error(),
            "we can't remove same peer twice");

    let change_peer = new_admin_request(1,
                                        new_change_peer_cmd(ConfChangeType::RemoveNode,
                                                            new_peer(3, 3, 3),
                                                            &stale_epoch));
    let resp = cluster.call_command_on_leader(1, change_peer, Duration::from_secs(3)).unwrap();
    assert!(resp.get_header().has_error(),
            "We can't change peer with stale epoch");

    // peer 3 must exist
    must_get_equal(&engine_3, b"a3", b"v3");

    // add peer 2 then remove it again.
    cluster.change_peer(r1, ConfChangeType::AddNode, new_peer(2, 2, 2));

    // Force update a2 to check whether peer 2 added ok and received the snapshot.
    let (key, value) = (b"a2", b"v2");
    cluster.put(key, value);
    assert_eq!(cluster.get(key), Some(value.to_vec()));

    let engine_2 = cluster.get_engine(2);

    must_get_equal(&engine_2, b"a1", b"v1");
    must_get_equal(&engine_2, b"a2", b"v2");
    must_get_equal(&engine_2, b"a3", b"v3");

    // Remove peer (2, 2, 2) from region 1.
    cluster.change_peer(r1, ConfChangeType::RemoveNode, new_peer(2, 2, 2));

    // add peer (2, 2, 4) to region 1.
    cluster.change_peer(r1, ConfChangeType::AddNode, new_peer(2, 2, 4));
    // Remove peer (3, 3, 3) from region 1.
    cluster.change_peer(r1, ConfChangeType::RemoveNode, new_peer(3, 3, 3));

    let (key, value) = (b"a4", b"v4");
    cluster.put(key, value);
    assert_eq!(cluster.get(key), Some(value.to_vec()));
    // now peer 4 in store 2 must have v1, v2, v3, v4, we check v1 and v4 here.
    let engine_2 = cluster.get_engine(2);

    must_get_equal(&engine_2, b"a1", b"v1");
    must_get_equal(&engine_2, b"a4", b"v4");

    // peer 3 has nothing, we check v1 and v4 here.
    must_get_none(&engine_3, b"a1");
    must_get_none(&engine_3, b"a4");

    // TODO: add more tests.
}

fn new_conf_change_peer(store: &metapb::Store,
                        pd_client: &Arc<RwLock<TestPdClient>>)
                        -> metapb::Peer {
    let peer_id = pd_client.wl().alloc_id().unwrap();
    new_peer(store.get_node_id(), store.get_store_id(), peer_id)
}

fn test_pd_conf_change<T: Simulator>(cluster: &mut Cluster<T>) {
    // init_env_log();
    cluster.start();

    let cluster_id = cluster.id();
    let pd_client = cluster.pd_client.clone();
    let region = &pd_client.rl().get_region(cluster_id, b"").unwrap();
    let region_id = region.get_region_id();

    let mut stores = pd_client.rl().get_stores(cluster_id).unwrap();

    // Must have only one peer
    assert_eq!(region.get_peers().len(), 1);

    let peer = &region.get_peers()[0];

    let i = stores.iter().position(|store| store.get_store_id() == peer.get_store_id()).unwrap();
    stores.swap(0, i);

    // Now the first store has first region. others have none.

    let (key, value) = (b"a1", b"v1");
    cluster.put(key, value);
    assert_eq!(cluster.get(key), Some(value.to_vec()));

    let peer2 = new_conf_change_peer(&stores[1], &pd_client);
    let engine_2 = cluster.get_engine(peer2.get_node_id());
    assert!(engine_2.get_value(&keys::data_key(b"a1")).unwrap().is_none());
    // add new peer to first region.
    cluster.change_peer(region_id, ConfChangeType::AddNode, peer2.clone());

    let (key, value) = (b"a2", b"v2");
    cluster.put(key, value);
    assert_eq!(cluster.get(key), Some(value.to_vec()));

    // now peer 2 must have v1 and v2;
    must_get_equal(&engine_2, b"a1", b"v1");
    must_get_equal(&engine_2, b"a2", b"v2");

    // add new peer to first region.
    let peer3 = new_conf_change_peer(&stores[2], &pd_client);
    cluster.change_peer(region_id, ConfChangeType::AddNode, peer3.clone());
    // Remove peer2 from first region.
    cluster.change_peer(region_id, ConfChangeType::RemoveNode, peer2.clone());

    let (key, value) = (b"a3", b"v3");
    cluster.put(key, value);
    assert_eq!(cluster.get(key), Some(value.to_vec()));
    // now peer 3 must have v1, v2 and v3
    let engine_3 = cluster.get_engine(peer3.get_node_id());
    must_get_equal(&engine_3, b"a1", b"v1");
    must_get_equal(&engine_3, b"a2", b"v2");
    must_get_equal(&engine_3, b"a3", b"v3");

    // peer 2 has nothing
    must_get_none(&engine_2, b"a1");
    must_get_none(&engine_2, b"a2");
    // add peer4 to first region 1.
    let peer4 = new_conf_change_peer(&stores[1], &pd_client);
    cluster.change_peer(region_id, ConfChangeType::AddNode, peer4.clone());
    // Remove peer3 from first region.
    cluster.change_peer(region_id, ConfChangeType::RemoveNode, peer3.clone());

    let (key, value) = (b"a4", b"v4");
    cluster.put(key, value);
    assert_eq!(cluster.get(key), Some(value.to_vec()));
    // now peer4 must have v1, v2, v3, v4, we check v1 and v4 here.
    let engine_2 = cluster.get_engine(peer4.get_node_id());

    must_get_equal(&engine_2, b"a1", b"v1");
    must_get_equal(&engine_2, b"a4", b"v4");

    // peer 3 has nothing, we check v1 and v4 here.
    must_get_none(&engine_3, b"a1");
    must_get_none(&engine_3, b"a4");

    // TODO: add more tests.
}

#[test]
fn test_node_simple_conf_change() {
    let count = 5;
    let mut cluster = new_node_cluster(0, count);
    test_simple_conf_change(&mut cluster);
}

#[test]
fn test_server_simple_conf_change() {
    let count = 5;
    let mut cluster = new_server_cluster(0, count);
    test_simple_conf_change(&mut cluster);
}

#[test]
fn test_node_pd_conf_change() {
    let count = 5;
    let mut cluster = new_node_cluster(0, count);
    test_pd_conf_change(&mut cluster);
}

#[test]
fn test_server_pd_conf_change() {
    let count = 5;
    let mut cluster = new_server_cluster(0, count);
    test_pd_conf_change(&mut cluster);
}

fn test_auto_adjust_replica<T: Simulator>(cluster: &mut Cluster<T>) {
    cluster.cfg.store_cfg.replica_check_tick_interval = 200;
    cluster.start();

    let cluster_id = cluster.id();
    let pd_client = cluster.pd_client.clone();
    let region = pd_client.rl().get_region(cluster_id, b"").unwrap();
    let region_id = region.get_region_id();

    let stores = pd_client.rl().get_stores(cluster_id).unwrap();

    // Must have only one peer
    assert_eq!(region.get_peers().len(), 1);

    thread::sleep(Duration::from_secs(1));

    let region = pd_client.rl().get_region_by_id(cluster_id, region_id).unwrap();
    // default replica is 5.
    assert_eq!(region.get_peers().len(), 5);

    let (key, value) = (b"a1", b"v1");
    cluster.put(key, value);
    assert_eq!(cluster.get(key), Some(value.to_vec()));

    let i = stores.iter()
                  .position(|s| {
                      region.get_peers().iter().all(|p| s.get_store_id() != p.get_store_id())
                  })
                  .unwrap();

    let peer = new_conf_change_peer(&stores[i], &pd_client);
    let engine = cluster.get_engine(peer.get_node_id());
    must_get_none(&engine, b"a1");

    cluster.change_peer(region_id, ConfChangeType::AddNode, peer.clone());

    let region = pd_client.rl().get_region_by_id(cluster_id, region_id).unwrap();
    assert_eq!(region.get_peers().len(), 6);

    thread::sleep(Duration::from_millis(300));

    let region = pd_client.rl().get_region_by_id(cluster_id, region_id).unwrap();
    assert_eq!(region.get_peers().len(), 5);

    let peer = region.get_peers().get(1).unwrap().clone();

    cluster.change_peer(region_id, ConfChangeType::RemoveNode, peer);
    let region = pd_client.rl().get_region_by_id(cluster_id, region_id).unwrap();
    assert_eq!(region.get_peers().len(), 4);

    thread::sleep(Duration::from_millis(300));

    let region = pd_client.rl().get_region_by_id(cluster_id, region_id).unwrap();
    assert_eq!(region.get_peers().len(), 5);
}

#[test]
fn test_node_auto_adjust_replica() {
    let count = 7;
    let mut cluster = new_node_cluster(0, count);
    test_auto_adjust_replica(&mut cluster);
}

#[test]
fn test_server_auto_adjust_replica() {
    let count = 7;
    let mut cluster = new_server_cluster(0, count);
    test_auto_adjust_replica(&mut cluster);
}
<commit_msg>raftserver: check epoch conf version.<commit_after>use std::time::Duration;
use std::sync::{Arc, RwLock};
use std::thread;

use tikv::raftserver::store::*;
use kvproto::raftpb::ConfChangeType;
use kvproto::metapb;
use tikv::pd::PdClient;
use tikv::util::HandyRwLock;

use super::cluster::{Cluster, Simulator};
use super::node::new_node_cluster;
use super::server::new_server_cluster;
use super::util::*;
use super::pd::TestPdClient;

fn test_simple_conf_change<T: Simulator>(cluster: &mut Cluster<T>) {
    // init_env_log();
    let r1 = cluster.bootstrap_conf_change();

    cluster.start();

    // Now region 1 only has peer (1, 1, 1);
    let (key, value) = (b"a1", b"v1");

    cluster.put(key, value);
    assert_eq!(cluster.get(key), Some(value.to_vec()));

    let engine_2 = cluster.get_engine(2);
    must_get_none(&engine_2, b"a1");
    // add peer (2,2,2) to region 1.
    cluster.change_peer(r1, ConfChangeType::AddNode, new_peer(2, 2, 2));

    let (key, value) = (b"a2", b"v2");
    cluster.put(key, value);
    assert_eq!(cluster.get(key), Some(value.to_vec()));

    // now peer 2 must have v1 and v2;
    must_get_equal(&engine_2, b"a1", b"v1");
    must_get_equal(&engine_2, b"a2", b"v2");

    let epoch = cluster.pd_client
                       .rl()
                       .get_region_by_id(cluster.id(), 1)
                       .unwrap()
                       .get_region_epoch()
                       .clone();

    // Conf version must change.
    assert!(epoch.get_conf_ver() > 1);

    let change_peer = new_admin_request(1,
                                        new_change_peer_cmd(ConfChangeType::AddNode,
                                                            new_peer(2, 2, 2),
                                                            &epoch));
    let resp = cluster.call_command_on_leader(1, change_peer, Duration::from_secs(3)).unwrap();
    assert!(resp.get_header().has_error(),
            "we can't add same peer twice");

    // Send an invalid stale epoch
    let mut stale_epoch = metapb::RegionEpoch::new();
    stale_epoch.set_version(1);
    stale_epoch.set_conf_ver(1);
    let change_peer = new_admin_request(1,
                                        new_change_peer_cmd(ConfChangeType::AddNode,
                                                            new_peer(5, 5, 5),
                                                            &stale_epoch));
    let resp = cluster.call_command_on_leader(1, change_peer, Duration::from_secs(3)).unwrap();
    assert!(resp.get_header().has_error(),
            "We can't change peer with stale epoch");

    // peer 5 must not exist
    let engine_5 = cluster.get_engine(5);
    must_get_none(&engine_5, b"a1");

    // add peer (3, 3, 3) to region 1.
    cluster.change_peer(r1, ConfChangeType::AddNode, new_peer(3, 3, 3));
    // Remove peer (2, 2, 2) from region 1.
    cluster.change_peer(r1, ConfChangeType::RemoveNode, new_peer(2, 2, 2));

    let (key, value) = (b"a3", b"v3");
    cluster.put(key, value);
    assert_eq!(cluster.get(key), Some(value.to_vec()));
    // now peer 3 must have v1, v2 and v3
    let engine_3 = cluster.get_engine(3);
    must_get_equal(&engine_3, b"a1", b"v1");
    must_get_equal(&engine_3, b"a2", b"v2");
    must_get_equal(&engine_3, b"a3", b"v3");

    // peer 2 has nothing
    must_get_none(&engine_2, b"a1");
    must_get_none(&engine_2, b"a2");

    let epoch = cluster.pd_client
                       .rl()
                       .get_region_by_id(cluster.id(), 1)
                       .unwrap()
                       .get_region_epoch()
                       .clone();
    let change_peer = new_admin_request(1,
                                        new_change_peer_cmd(ConfChangeType::RemoveNode,
                                                            new_peer(2, 2, 2),
                                                            &epoch));
    let resp = cluster.call_command_on_leader(1, change_peer, Duration::from_secs(3)).unwrap();
    assert!(resp.get_header().has_error(),
            "we can't remove same peer twice");

    let change_peer = new_admin_request(1,
                                        new_change_peer_cmd(ConfChangeType::RemoveNode,
                                                            new_peer(3, 3, 3),
                                                            &stale_epoch));
    let resp = cluster.call_command_on_leader(1, change_peer, Duration::from_secs(3)).unwrap();
    assert!(resp.get_header().has_error(),
            "We can't change peer with stale epoch");

    // peer 3 must exist
    must_get_equal(&engine_3, b"a3", b"v3");

    // add peer 2 then remove it again.
    cluster.change_peer(r1, ConfChangeType::AddNode, new_peer(2, 2, 2));

    // Force update a2 to check whether peer 2 added ok and received the snapshot.
    let (key, value) = (b"a2", b"v2");
    cluster.put(key, value);
    assert_eq!(cluster.get(key), Some(value.to_vec()));

    let engine_2 = cluster.get_engine(2);

    must_get_equal(&engine_2, b"a1", b"v1");
    must_get_equal(&engine_2, b"a2", b"v2");
    must_get_equal(&engine_2, b"a3", b"v3");

    // Remove peer (2, 2, 2) from region 1.
    cluster.change_peer(r1, ConfChangeType::RemoveNode, new_peer(2, 2, 2));

    // add peer (2, 2, 4) to region 1.
    cluster.change_peer(r1, ConfChangeType::AddNode, new_peer(2, 2, 4));
    // Remove peer (3, 3, 3) from region 1.
    cluster.change_peer(r1, ConfChangeType::RemoveNode, new_peer(3, 3, 3));

    let (key, value) = (b"a4", b"v4");
    cluster.put(key, value);
    assert_eq!(cluster.get(key), Some(value.to_vec()));
    // now peer 4 in store 2 must have v1, v2, v3, v4, we check v1 and v4 here.
    let engine_2 = cluster.get_engine(2);

    must_get_equal(&engine_2, b"a1", b"v1");
    must_get_equal(&engine_2, b"a4", b"v4");

    // peer 3 has nothing, we check v1 and v4 here.
    must_get_none(&engine_3, b"a1");
    must_get_none(&engine_3, b"a4");

    // TODO: add more tests.
}

fn new_conf_change_peer(store: &metapb::Store,
                        pd_client: &Arc<RwLock<TestPdClient>>)
                        -> metapb::Peer {
    let peer_id = pd_client.wl().alloc_id().unwrap();
    new_peer(store.get_node_id(), store.get_store_id(), peer_id)
}

fn test_pd_conf_change<T: Simulator>(cluster: &mut Cluster<T>) {
    // init_env_log();
    cluster.start();

    let cluster_id = cluster.id();
    let pd_client = cluster.pd_client.clone();
    let region = &pd_client.rl().get_region(cluster_id, b"").unwrap();
    let region_id = region.get_region_id();

    let mut stores = pd_client.rl().get_stores(cluster_id).unwrap();

    // Must have only one peer
    assert_eq!(region.get_peers().len(), 1);

    let peer = &region.get_peers()[0];

    let i = stores.iter().position(|store| store.get_store_id() == peer.get_store_id()).unwrap();
    stores.swap(0, i);

    // Now the first store has first region. others have none.

    let (key, value) = (b"a1", b"v1");
    cluster.put(key, value);
    assert_eq!(cluster.get(key), Some(value.to_vec()));

    let peer2 = new_conf_change_peer(&stores[1], &pd_client);
    let engine_2 = cluster.get_engine(peer2.get_node_id());
    assert!(engine_2.get_value(&keys::data_key(b"a1")).unwrap().is_none());
    // add new peer to first region.
    cluster.change_peer(region_id, ConfChangeType::AddNode, peer2.clone());

    let (key, value) = (b"a2", b"v2");
    cluster.put(key, value);
    assert_eq!(cluster.get(key), Some(value.to_vec()));

    // now peer 2 must have v1 and v2;
    must_get_equal(&engine_2, b"a1", b"v1");
    must_get_equal(&engine_2, b"a2", b"v2");

    // add new peer to first region.
    let peer3 = new_conf_change_peer(&stores[2], &pd_client);
    cluster.change_peer(region_id, ConfChangeType::AddNode, peer3.clone());
    // Remove peer2 from first region.
    cluster.change_peer(region_id, ConfChangeType::RemoveNode, peer2.clone());

    let (key, value) = (b"a3", b"v3");
    cluster.put(key, value);
    assert_eq!(cluster.get(key), Some(value.to_vec()));
    // now peer 3 must have v1, v2 and v3
    let engine_3 = cluster.get_engine(peer3.get_node_id());
    must_get_equal(&engine_3, b"a1", b"v1");
    must_get_equal(&engine_3, b"a2", b"v2");
    must_get_equal(&engine_3, b"a3", b"v3");

    // peer 2 has nothing
    must_get_none(&engine_2, b"a1");
    must_get_none(&engine_2, b"a2");
    // add peer4 to first region 1.
    let peer4 = new_conf_change_peer(&stores[1], &pd_client);
    cluster.change_peer(region_id, ConfChangeType::AddNode, peer4.clone());
    // Remove peer3 from first region.
    cluster.change_peer(region_id, ConfChangeType::RemoveNode, peer3.clone());

    let (key, value) = (b"a4", b"v4");
    cluster.put(key, value);
    assert_eq!(cluster.get(key), Some(value.to_vec()));
    // now peer4 must have v1, v2, v3, v4, we check v1 and v4 here.
    let engine_2 = cluster.get_engine(peer4.get_node_id());

    must_get_equal(&engine_2, b"a1", b"v1");
    must_get_equal(&engine_2, b"a4", b"v4");

    // peer 3 has nothing, we check v1 and v4 here.
    must_get_none(&engine_3, b"a1");
    must_get_none(&engine_3, b"a4");

    // TODO: add more tests.
}

#[test]
fn test_node_simple_conf_change() {
    let count = 5;
    let mut cluster = new_node_cluster(0, count);
    test_simple_conf_change(&mut cluster);
}

#[test]
fn test_server_simple_conf_change() {
    let count = 5;
    let mut cluster = new_server_cluster(0, count);
    test_simple_conf_change(&mut cluster);
}

#[test]
fn test_node_pd_conf_change() {
    let count = 5;
    let mut cluster = new_node_cluster(0, count);
    test_pd_conf_change(&mut cluster);
}

#[test]
fn test_server_pd_conf_change() {
    let count = 5;
    let mut cluster = new_server_cluster(0, count);
    test_pd_conf_change(&mut cluster);
}

fn test_auto_adjust_replica<T: Simulator>(cluster: &mut Cluster<T>) {
    cluster.cfg.store_cfg.replica_check_tick_interval = 200;
    cluster.start();

    let cluster_id = cluster.id();
    let pd_client = cluster.pd_client.clone();
    let region = pd_client.rl().get_region(cluster_id, b"").unwrap();
    let region_id = region.get_region_id();

    let stores = pd_client.rl().get_stores(cluster_id).unwrap();

    // Must have only one peer
    assert_eq!(region.get_peers().len(), 1);

    thread::sleep(Duration::from_secs(1));

    let region = pd_client.rl().get_region_by_id(cluster_id, region_id).unwrap();
    // default replica is 5.
    assert_eq!(region.get_peers().len(), 5);

    let (key, value) = (b"a1", b"v1");
    cluster.put(key, value);
    assert_eq!(cluster.get(key), Some(value.to_vec()));

    let i = stores.iter()
                  .position(|s| {
                      region.get_peers().iter().all(|p| s.get_store_id() != p.get_store_id())
                  })
                  .unwrap();

    let peer = new_conf_change_peer(&stores[i], &pd_client);
    let engine = cluster.get_engine(peer.get_node_id());
    must_get_none(&engine, b"a1");

    cluster.change_peer(region_id, ConfChangeType::AddNode, peer.clone());

    let region = pd_client.rl().get_region_by_id(cluster_id, region_id).unwrap();
    assert_eq!(region.get_peers().len(), 6);

    thread::sleep(Duration::from_millis(300));

    let region = pd_client.rl().get_region_by_id(cluster_id, region_id).unwrap();
    assert_eq!(region.get_peers().len(), 5);

    let peer = region.get_peers().get(1).unwrap().clone();

    cluster.change_peer(region_id, ConfChangeType::RemoveNode, peer);
    let region = pd_client.rl().get_region_by_id(cluster_id, region_id).unwrap();
    assert_eq!(region.get_peers().len(), 4);

    thread::sleep(Duration::from_millis(300));

    let region = pd_client.rl().get_region_by_id(cluster_id, region_id).unwrap();
    assert_eq!(region.get_peers().len(), 5);
}

#[test]
fn test_node_auto_adjust_replica() {
    let count = 7;
    let mut cluster = new_node_cluster(0, count);
    test_auto_adjust_replica(&mut cluster);
}

#[test]
fn test_server_auto_adjust_replica() {
    let count = 7;
    let mut cluster = new_server_cluster(0, count);
    test_auto_adjust_replica(&mut cluster);
}
<|endoftext|>


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

//! `OpCounters` is a collection of convenience methods to add arbitrary counters to modules.
//! For now, it supports Int-Counters, Int-Gauges, and Histogram.

use prometheus::{
    core::{Collector, Desc},
    proto::MetricFamily,
    Histogram, HistogramOpts, HistogramTimer, HistogramVec, IntCounterVec, IntGauge, IntGaugeVec,
    Opts,
};

use std::time::Duration;

/// A small wrapper around Histogram to handle the special case
/// of duration buckets.
/// This Histogram will handle the correct granularity for logging
/// time duration in a way that fits the used buckets.
pub struct DurationHistogram {
    histogram: Histogram,
}

impl DurationHistogram {
    pub fn new(histogram: Histogram) -> DurationHistogram {
        DurationHistogram { histogram }
    }

    pub fn observe_duration(&self, d: Duration) {
        // Duration is full seconds + nanos elapsed from the previous full second
        let v = d.as_secs() as f64 + f64::from(d.subsec_nanos()) / 1e9;
        self.histogram.observe(v);
    }
}

#[derive(Clone)]
pub struct OpMetrics {
    module: String,
    counters: IntCounterVec,
    gauges: IntGaugeVec,
    peer_gauges: IntGaugeVec,
    duration_histograms: HistogramVec,
}

impl OpMetrics {
    pub fn new<S: Into<String>>(name: S) -> OpMetrics {
        let name_str = name.into();
        OpMetrics {
            module: name_str.clone(),
            counters: IntCounterVec::new(
                Opts::new(name_str.clone(), format!("Counters for {}", name_str)),
                &["op"],
            )
            .unwrap(),
            gauges: IntGaugeVec::new(
                Opts::new(
                    format!("{}_gauge", name_str.clone()),
                    format!("Gauges for {}", name_str),
                ),
                &["op"],
            )
            .unwrap(),
            peer_gauges: IntGaugeVec::new(
                Opts::new(
                    format!("{}_peer_gauge", name_str.clone()),
                    format!("Gauges of each remote peer for {}", name_str),
                ),
                &["op", "remote_peer_id"],
            )
            .unwrap(),
            duration_histograms: HistogramVec::new(
                HistogramOpts::new(
                    format!("{}_duration", name_str.clone()),
                    format!("Histogram values for {}", name_str),
                ),
                &["op"],
            )
            .unwrap(),
        }
    }

    pub fn new_and_registered<S: Into<String>>(name: S) -> OpMetrics {
        let op_metrics = OpMetrics::new(name);
        prometheus::register(Box::new(op_metrics.clone()))
            .expect("OpMetrics registration on Prometheus failed.");
        op_metrics
    }

    #[inline]
    pub fn gauge(&self, name: &str) -> IntGauge {
        self.gauges.with_label_values(&[name])
    }

    #[inline]
    pub fn peer_gauge(&self, name: &str, remote_peer_id: &str) -> IntGauge {
        self.peer_gauges.with_label_values(&[name, remote_peer_id])
    }

    pub fn duration_histogram(&self, name: &str) -> DurationHistogram {
        DurationHistogram::new(self.duration_histograms.with_label_values(&[name]))
    }

    #[inline]
    pub fn inc(&self, op: &str) {
        self.counters.with_label_values(&[op]).inc();
    }

    #[inline]
    pub fn inc_by(&self, op: &str, v: usize) {
        // The underlying method is expecting i64, but most of the types
        // we're going to log are `u64` or `usize`.
        self.counters.with_label_values(&[op]).inc_by(v as i64);
    }

    #[inline]
    pub fn add(&self, op: &str) {
        self.gauges.with_label_values(&[op]).inc();
    }

    #[inline]
    pub fn sub(&self, op: &str) {
        self.gauges.with_label_values(&[op]).dec();
    }

    #[inline]
    pub fn set(&self, op: &str, v: usize) {
        // The underlying method is expecting i64, but most of the types
        // we're going to log are `u64` or `usize`.
        self.gauges.with_label_values(&[op]).set(v as i64);
    }

    #[inline]
    pub fn observe(&self, op: &str, v: f64) {
        self.duration_histograms.with_label_values(&[op]).observe(v);
    }

    pub fn observe_duration(&self, op: &str, d: Duration) {
        // Duration is full seconds + nanos elapsed from the previous full second
        let v = d.as_secs() as f64 + f64::from(d.subsec_nanos()) / 1e9;
        self.duration_histograms.with_label_values(&[op]).observe(v);
    }

    pub fn timer(&self, op: &str) -> HistogramTimer {
        self.duration_histograms
            .with_label_values(&[op])
            .start_timer()
    }
}

impl Collector for OpMetrics {
    fn desc(&self) -> Vec<&Desc> {
        let mut ms = Vec::with_capacity(4);
        ms.extend(self.counters.desc());
        ms.extend(self.gauges.desc());
        ms.extend(self.peer_gauges.desc());
        ms.extend(self.duration_histograms.desc());
        ms
    }

    fn collect(&self) -> Vec<MetricFamily> {
        let mut ms = Vec::with_capacity(4);
        ms.extend(self.counters.collect());
        ms.extend(self.gauges.collect());
        ms.extend(self.peer_gauges.collect());
        ms.extend(self.duration_histograms.collect());
        ms
    }
}
<commit_msg>Removing duration_histogram from common/metrics/src/op_counters.rs<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! `OpCounters` is a collection of convenience methods to add arbitrary counters to modules.
//! For now, it supports Int-Counters, Int-Gauges, and Histogram.

use prometheus::{
    core::{Collector, Desc},
    proto::MetricFamily,
    Histogram, HistogramOpts, HistogramTimer, HistogramVec, IntCounterVec, IntGauge, IntGaugeVec,
    Opts,
};

use std::time::Duration;

/// A small wrapper around Histogram to handle the special case
/// of duration buckets.
/// This Histogram will handle the correct granularity for logging
/// time duration in a way that fits the used buckets.
pub struct DurationHistogram {
    histogram: Histogram,
}

impl DurationHistogram {
    pub fn new(histogram: Histogram) -> DurationHistogram {
        DurationHistogram { histogram }
    }

    pub fn observe_duration(&self, d: Duration) {
        // Duration is full seconds + nanos elapsed from the previous full second
        let v = d.as_secs() as f64 + f64::from(d.subsec_nanos()) / 1e9;
        self.histogram.observe(v);
    }
}

#[derive(Clone)]
pub struct OpMetrics {
    module: String,
    counters: IntCounterVec,
    gauges: IntGaugeVec,
    peer_gauges: IntGaugeVec,
    duration_histograms: HistogramVec,
}

impl OpMetrics {
    pub fn new<S: Into<String>>(name: S) -> OpMetrics {
        let name_str = name.into();
        OpMetrics {
            module: name_str.clone(),
            counters: IntCounterVec::new(
                Opts::new(name_str.clone(), format!("Counters for {}", name_str)),
                &["op"],
            )
            .unwrap(),
            gauges: IntGaugeVec::new(
                Opts::new(
                    format!("{}_gauge", name_str.clone()),
                    format!("Gauges for {}", name_str),
                ),
                &["op"],
            )
            .unwrap(),
            peer_gauges: IntGaugeVec::new(
                Opts::new(
                    format!("{}_peer_gauge", name_str.clone()),
                    format!("Gauges of each remote peer for {}", name_str),
                ),
                &["op", "remote_peer_id"],
            )
            .unwrap(),
            duration_histograms: HistogramVec::new(
                HistogramOpts::new(
                    format!("{}_duration", name_str.clone()),
                    format!("Histogram values for {}", name_str),
                ),
                &["op"],
            )
            .unwrap(),
        }
    }

    pub fn new_and_registered<S: Into<String>>(name: S) -> OpMetrics {
        let op_metrics = OpMetrics::new(name);
        prometheus::register(Box::new(op_metrics.clone()))
            .expect("OpMetrics registration on Prometheus failed.");
        op_metrics
    }

    #[inline]
    pub fn gauge(&self, name: &str) -> IntGauge {
        self.gauges.with_label_values(&[name])
    }

    #[inline]
    pub fn peer_gauge(&self, name: &str, remote_peer_id: &str) -> IntGauge {
        self.peer_gauges.with_label_values(&[name, remote_peer_id])
    }

    #[inline]
    pub fn inc(&self, op: &str) {
        self.counters.with_label_values(&[op]).inc();
    }

    #[inline]
    pub fn inc_by(&self, op: &str, v: usize) {
        // The underlying method is expecting i64, but most of the types
        // we're going to log are `u64` or `usize`.
        self.counters.with_label_values(&[op]).inc_by(v as i64);
    }

    #[inline]
    pub fn add(&self, op: &str) {
        self.gauges.with_label_values(&[op]).inc();
    }

    #[inline]
    pub fn sub(&self, op: &str) {
        self.gauges.with_label_values(&[op]).dec();
    }

    #[inline]
    pub fn set(&self, op: &str, v: usize) {
        // The underlying method is expecting i64, but most of the types
        // we're going to log are `u64` or `usize`.
        self.gauges.with_label_values(&[op]).set(v as i64);
    }

    #[inline]
    pub fn observe(&self, op: &str, v: f64) {
        self.duration_histograms.with_label_values(&[op]).observe(v);
    }

    pub fn observe_duration(&self, op: &str, d: Duration) {
        // Duration is full seconds + nanos elapsed from the previous full second
        let v = d.as_secs() as f64 + f64::from(d.subsec_nanos()) / 1e9;
        self.duration_histograms.with_label_values(&[op]).observe(v);
    }

    pub fn timer(&self, op: &str) -> HistogramTimer {
        self.duration_histograms
            .with_label_values(&[op])
            .start_timer()
    }
}

impl Collector for OpMetrics {
    fn desc(&self) -> Vec<&Desc> {
        let mut ms = Vec::with_capacity(4);
        ms.extend(self.counters.desc());
        ms.extend(self.gauges.desc());
        ms.extend(self.peer_gauges.desc());
        ms.extend(self.duration_histograms.desc());
        ms
    }

    fn collect(&self) -> Vec<MetricFamily> {
        let mut ms = Vec::with_capacity(4);
        ms.extend(self.counters.collect());
        ms.extend(self.gauges.collect());
        ms.extend(self.peer_gauges.collect());
        ms.extend(self.duration_histograms.collect());
        ms
    }
}
<|endoftext|>


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

use dom::bindings::js::{JS, Root};
use dom::bindings::trace::JSTraceable;
use dom::globalscope::GlobalScope;
use js::jsapi::JSTracer;
use std::cell::RefCell;

thread_local!(static STACK: RefCell<Vec<StackEntry>> = RefCell::new(Vec::new()));

#[allow(unrooted_must_root)]
#[derive(JSTraceable)]
struct StackEntry {
    global: JS<GlobalScope>,
}

/// Traces the script settings stack.
pub unsafe fn trace(tracer: *mut JSTracer) {
    STACK.with(|stack| {
        stack.borrow().trace(tracer);
    })
}

/// RAII struct that pushes and pops entries from the script settings stack.
pub struct AutoEntryScript {
    global: *const GlobalScope,
}

impl AutoEntryScript {
    /// https://html.spec.whatwg.org/multipage/#prepare-to-run-script
    pub fn new(global: &GlobalScope) -> Self {
        STACK.with(|stack| {
            trace!("Prepare to run script with {:p}", global);
            let mut stack = stack.borrow_mut();
            stack.push(StackEntry {
                global: JS::from_ref(global),
            });
            AutoEntryScript {
                global: global as *const _,
            }
        })
    }
}

impl Drop for AutoEntryScript {
    /// https://html.spec.whatwg.org/multipage/#clean-up-after-running-script
    fn drop(&mut self) {
        STACK.with(|stack| {
            let mut stack = stack.borrow_mut();
            let entry = stack.pop().unwrap();
            assert_eq!(&*entry.global as *const GlobalScope,
                       self.global,
                       "Dropped AutoEntryScript out of order.");
            trace!("Clean up after running script with {:p}", self.global);
        })
    }
}

/// Returns the ["entry"] global object.
///
/// ["entry"]: https://html.spec.whatwg.org/multipage/#entry
pub fn entry_global() -> Root<GlobalScope> {
    STACK.with(|stack| {
        stack.borrow()
             .last()
             .map(|entry| Root::from_ref(&*entry.global))
    }).unwrap()
}
<commit_msg>Use usize for AutoEntryScript::global.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use dom::bindings::js::{JS, Root};
use dom::bindings::trace::JSTraceable;
use dom::globalscope::GlobalScope;
use js::jsapi::JSTracer;
use std::cell::RefCell;

thread_local!(static STACK: RefCell<Vec<StackEntry>> = RefCell::new(Vec::new()));

#[allow(unrooted_must_root)]
#[derive(JSTraceable)]
struct StackEntry {
    global: JS<GlobalScope>,
}

/// Traces the script settings stack.
pub unsafe fn trace(tracer: *mut JSTracer) {
    STACK.with(|stack| {
        stack.borrow().trace(tracer);
    })
}

/// RAII struct that pushes and pops entries from the script settings stack.
pub struct AutoEntryScript {
    global: usize,
}

impl AutoEntryScript {
    /// https://html.spec.whatwg.org/multipage/#prepare-to-run-script
    pub fn new(global: &GlobalScope) -> Self {
        STACK.with(|stack| {
            trace!("Prepare to run script with {:p}", global);
            let mut stack = stack.borrow_mut();
            stack.push(StackEntry {
                global: JS::from_ref(global),
            });
            AutoEntryScript {
                global: global as *const _ as usize,
            }
        })
    }
}

impl Drop for AutoEntryScript {
    /// https://html.spec.whatwg.org/multipage/#clean-up-after-running-script
    fn drop(&mut self) {
        STACK.with(|stack| {
            let mut stack = stack.borrow_mut();
            let entry = stack.pop().unwrap();
            assert_eq!(&*entry.global as *const GlobalScope as usize,
                       self.global,
                       "Dropped AutoEntryScript out of order.");
            trace!("Clean up after running script with {:p}", &*entry.global);
        })
    }
}

/// Returns the ["entry"] global object.
///
/// ["entry"]: https://html.spec.whatwg.org/multipage/#entry
pub fn entry_global() -> Root<GlobalScope> {
    STACK.with(|stack| {
        stack.borrow()
             .last()
             .map(|entry| Root::from_ref(&*entry.global))
    }).unwrap()
}
<|endoftext|>


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

//! Defines aggregations over Arrow arrays.

use std::ops::Add;

use crate::array::{
    Array, LargeStringArray, PrimitiveArray, StringArray, StringArrayOps,
};
use crate::datatypes::ArrowNumericType;

/// Helper macro to perform min/max of strings
macro_rules! min_max_string_helper {
    ($array:expr, $cmp:tt) => {{
        let null_count = $array.null_count();

        if null_count == $array.len() {
            return None
        }
        let mut n = "";
        let mut has_value = false;
        let data = $array.data();

        if null_count == 0 {
            for i in 0..data.len() {
                let item = $array.value(i);
                if !has_value || (&n $cmp &item) {
                    has_value = true;
                    n = item;
                }
            }
        } else {
            for i in 0..data.len() {
                let item = $array.value(i);
                if !has_value || data.is_valid(i) && (&n $cmp &item) {
                    has_value = true;
                    n = item;
                }
            }
        }
        Some(n)
    }}
}

/// Returns the minimum value in the array, according to the natural order.
pub fn min<T>(array: &PrimitiveArray<T>) -> Option<T::Native>
where
    T: ArrowNumericType,
{
    min_max_helper(array, |a, b| a > b)
}

/// Returns the maximum value in the array, according to the natural order.
pub fn max<T>(array: &PrimitiveArray<T>) -> Option<T::Native>
where
    T: ArrowNumericType,
{
    min_max_helper(array, |a, b| a < b)
}

/// Returns the maximum value in the string array, according to the natural order.
pub fn max_string(array: &StringArray) -> Option<&str> {
    min_max_string_helper!(array, <)
}

/// Returns the minimum value in the string array, according to the natural order.
pub fn min_string(array: &StringArray) -> Option<&str> {
    min_max_string_helper!(array, >)
}

/// Returns the minimum value in the string array, according to the natural order.
pub fn max_large_string(array: &LargeStringArray) -> Option<&str> {
    min_max_string_helper!(array, <)
}

/// Returns the minimum value in the string array, according to the natural order.
pub fn min_large_string(array: &LargeStringArray) -> Option<&str> {
    min_max_string_helper!(array, >)
}

/// Helper function to perform min/max lambda function on values from a numeric array.
fn min_max_helper<T, F>(array: &PrimitiveArray<T>, cmp: F) -> Option<T::Native>
where
    T: ArrowNumericType,
    F: Fn(&T::Native, &T::Native) -> bool,
{
    let null_count = array.null_count();

    if null_count == array.len() {
        return None;
    }

    let mut n: T::Native = T::default_value();
    let mut has_value = false;
    let data = array.data();
    let m = array.value_slice(0, data.len());

    if null_count == 0 {
        // optimized path for arrays without null values
        for item in m {
            if !has_value || cmp(&n, item) {
                has_value = true;
                n = *item
            }
        }
    } else {
        for (i, item) in m.iter().enumerate() {
            if !has_value || data.is_valid(i) && cmp(&n, item) {
                has_value = true;
                n = *item
            }
        }
    }
    Some(n)
}

/// Returns the sum of values in the array.
///
/// Returns `None` if the array is empty or only contains null values.
pub fn sum<T>(array: &PrimitiveArray<T>) -> Option<T::Native>
where
    T: ArrowNumericType,
    T::Native: Add<Output = T::Native>,
{
    let null_count = array.null_count();

    if null_count == array.len() {
        return None;
    }

    let mut n: T::Native = T::default_value();
    let data = array.data();
    let m = array.value_slice(0, data.len());

    if null_count == 0 {
        // optimized path for arrays without null values
        for item in m.iter().take(data.len()) {
            n = n + *item;
        }
    } else {
        for (i, item) in m.iter().enumerate() {
            if data.is_valid(i) {
                n = n + *item;
            }
        }
    }
    Some(n)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::array::*;

    #[test]
    fn test_primitive_array_sum() {
        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
        assert_eq!(15, sum(&a).unwrap());
    }

    #[test]
    fn test_primitive_array_float_sum() {
        let a = Float64Array::from(vec![1.1, 2.2, 3.3, 4.4, 5.5]);
        assert!(16.5 - sum(&a).unwrap() < f64::EPSILON);
    }

    #[test]
    fn test_primitive_array_sum_with_nulls() {
        let a = Int32Array::from(vec![None, Some(2), Some(3), None, Some(5)]);
        assert_eq!(10, sum(&a).unwrap());
    }

    #[test]
    fn test_primitive_array_sum_all_nulls() {
        let a = Int32Array::from(vec![None, None, None]);
        assert_eq!(None, sum(&a));
    }

    #[test]
    fn test_buffer_array_min_max() {
        let a = Int32Array::from(vec![5, 6, 7, 8, 9]);
        assert_eq!(5, min(&a).unwrap());
        assert_eq!(9, max(&a).unwrap());
    }

    #[test]
    fn test_buffer_array_min_max_with_nulls() {
        let a = Int32Array::from(vec![Some(5), None, None, Some(8), Some(9)]);
        assert_eq!(5, min(&a).unwrap());
        assert_eq!(9, max(&a).unwrap());
    }

    #[test]
    fn test_string_min_max_with_nulls() {
        let a = StringArray::from(vec![Some("b"), None, None, Some("a"), Some("c")]);
        assert_eq!("a", min_string(&a).unwrap());
        assert_eq!("c", max_string(&a).unwrap());
    }

    #[test]
    fn test_string_min_max_all_nulls() {
        let a = StringArray::from(vec![None, None]);
        assert_eq!(None, min_string(&a));
        assert_eq!(None, max_string(&a));
    }
}
<commit_msg>ARROW-10048: [Rust] Fixed error in computing min/max with null entries.<commit_after>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Defines aggregations over Arrow arrays.

use std::ops::Add;

use crate::array::{
    Array, LargeStringArray, PrimitiveArray, StringArray, StringArrayOps,
};
use crate::datatypes::ArrowNumericType;

/// Helper macro to perform min/max of strings
macro_rules! min_max_string_helper {
    ($array:expr, $cmp:tt) => {{
        let null_count = $array.null_count();

        if null_count == $array.len() {
            return None
        }
        let mut n = "";
        let mut has_value = false;
        let data = $array.data();

        if null_count == 0 {
            for i in 0..data.len() {
                let item = $array.value(i);
                if !has_value || (&n $cmp &item) {
                    has_value = true;
                    n = item;
                }
            }
        } else {
            for i in 0..data.len() {
                let item = $array.value(i);
                if data.is_valid(i) && (!has_value || (&n $cmp &item)) {
                    has_value = true;
                    n = item;
                }
            }
        }
        Some(n)
    }}
}

/// Returns the minimum value in the array, according to the natural order.
pub fn min<T>(array: &PrimitiveArray<T>) -> Option<T::Native>
where
    T: ArrowNumericType,
{
    min_max_helper(array, |a, b| a > b)
}

/// Returns the maximum value in the array, according to the natural order.
pub fn max<T>(array: &PrimitiveArray<T>) -> Option<T::Native>
where
    T: ArrowNumericType,
{
    min_max_helper(array, |a, b| a < b)
}

/// Returns the maximum value in the string array, according to the natural order.
pub fn max_string(array: &StringArray) -> Option<&str> {
    min_max_string_helper!(array, <)
}

/// Returns the minimum value in the string array, according to the natural order.
pub fn min_string(array: &StringArray) -> Option<&str> {
    min_max_string_helper!(array, >)
}

/// Returns the minimum value in the string array, according to the natural order.
pub fn max_large_string(array: &LargeStringArray) -> Option<&str> {
    min_max_string_helper!(array, <)
}

/// Returns the minimum value in the string array, according to the natural order.
pub fn min_large_string(array: &LargeStringArray) -> Option<&str> {
    min_max_string_helper!(array, >)
}

/// Helper function to perform min/max lambda function on values from a numeric array.
fn min_max_helper<T, F>(array: &PrimitiveArray<T>, cmp: F) -> Option<T::Native>
where
    T: ArrowNumericType,
    F: Fn(&T::Native, &T::Native) -> bool,
{
    let null_count = array.null_count();

    if null_count == array.len() {
        return None;
    }

    let mut n: T::Native = T::default_value();
    let mut has_value = false;
    let data = array.data();
    let m = array.value_slice(0, data.len());

    if null_count == 0 {
        // optimized path for arrays without null values
        for item in m {
            if !has_value || cmp(&n, item) {
                has_value = true;
                n = *item
            }
        }
    } else {
        for (i, item) in m.iter().enumerate() {
            if data.is_valid(i) && (!has_value || cmp(&n, item)) {
                has_value = true;
                n = *item
            }
        }
    }
    Some(n)
}

/// Returns the sum of values in the array.
///
/// Returns `None` if the array is empty or only contains null values.
pub fn sum<T>(array: &PrimitiveArray<T>) -> Option<T::Native>
where
    T: ArrowNumericType,
    T::Native: Add<Output = T::Native>,
{
    let null_count = array.null_count();

    if null_count == array.len() {
        return None;
    }

    let mut n: T::Native = T::default_value();
    let data = array.data();
    let m = array.value_slice(0, data.len());

    if null_count == 0 {
        // optimized path for arrays without null values
        for item in m.iter().take(data.len()) {
            n = n + *item;
        }
    } else {
        for (i, item) in m.iter().enumerate() {
            if data.is_valid(i) {
                n = n + *item;
            }
        }
    }
    Some(n)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::array::*;

    #[test]
    fn test_primitive_array_sum() {
        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
        assert_eq!(15, sum(&a).unwrap());
    }

    #[test]
    fn test_primitive_array_float_sum() {
        let a = Float64Array::from(vec![1.1, 2.2, 3.3, 4.4, 5.5]);
        assert!(16.5 - sum(&a).unwrap() < f64::EPSILON);
    }

    #[test]
    fn test_primitive_array_sum_with_nulls() {
        let a = Int32Array::from(vec![None, Some(2), Some(3), None, Some(5)]);
        assert_eq!(10, sum(&a).unwrap());
    }

    #[test]
    fn test_primitive_array_sum_all_nulls() {
        let a = Int32Array::from(vec![None, None, None]);
        assert_eq!(None, sum(&a));
    }

    #[test]
    fn test_buffer_array_min_max() {
        let a = Int32Array::from(vec![5, 6, 7, 8, 9]);
        assert_eq!(5, min(&a).unwrap());
        assert_eq!(9, max(&a).unwrap());
    }

    #[test]
    fn test_buffer_array_min_max_with_nulls() {
        let a = Int32Array::from(vec![Some(5), None, None, Some(8), Some(9)]);
        assert_eq!(5, min(&a).unwrap());
        assert_eq!(9, max(&a).unwrap());
    }

    #[test]
    fn test_buffer_min_max_1() {
        let a = Int32Array::from(vec![None, None, Some(5), Some(2)]);
        assert_eq!(Some(2), min(&a));
        assert_eq!(Some(5), max(&a));
    }

    #[test]
    fn test_string_min_max_with_nulls() {
        let a = StringArray::from(vec![Some("b"), None, None, Some("a"), Some("c")]);
        assert_eq!("a", min_string(&a).unwrap());
        assert_eq!("c", max_string(&a).unwrap());
    }

    #[test]
    fn test_string_min_max_all_nulls() {
        let a = StringArray::from(vec![None, None]);
        assert_eq!(None, min_string(&a));
        assert_eq!(None, max_string(&a));
    }

    #[test]
    fn test_string_min_max_1() {
        let a = StringArray::from(vec![None, None, Some("b"), Some("a")]);
        assert_eq!(Some("a"), min_string(&a));
        assert_eq!(Some("b"), max_string(&a));
    }
}
<|endoftext|>


<commit_before>//! Completion for associated items in a trait implementation.
//!
//! This module adds the completion items related to implementing associated
//! items within a `impl Trait for Struct` block. The current context node
//! must be within either a `FN`, `TYPE_ALIAS`, or `CONST` node
//! and an direct child of an `IMPL`.
//!
//! # Examples
//!
//! Considering the following trait `impl`:
//!
//! ```ignore
//! trait SomeTrait {
//!     fn foo();
//! }
//!
//! impl SomeTrait for () {
//!     fn f<|>
//! }
//! ```
//!
//! may result in the completion of the following method:
//!
//! ```ignore
//! # trait SomeTrait {
//! #    fn foo();
//! # }
//!
//! impl SomeTrait for () {
//!     fn foo() {}<|>
//! }
//! ```

use assists::utils::get_missing_assoc_items;
use hir::{self, HasAttrs, HasSource};
use syntax::{
    ast::{self, edit, Impl},
    AstNode, SyntaxKind, SyntaxNode, TextRange, T,
};
use text_edit::TextEdit;

use crate::{
    completion::{
        CompletionContext, CompletionItem, CompletionItemKind, CompletionKind, Completions,
    },
    display::function_declaration,
};

pub(crate) fn complete_trait_impl(acc: &mut Completions, ctx: &CompletionContext) {
    if let Some((trigger, impl_def)) = completion_match(ctx) {
        match trigger.kind() {
            SyntaxKind::NAME_REF => get_missing_assoc_items(&ctx.sema, &impl_def)
                .into_iter()
                .for_each(|item| match item {
                    hir::AssocItem::Function(fn_item) => {
                        add_function_impl(&trigger, acc, ctx, fn_item)
                    }
                    hir::AssocItem::TypeAlias(type_item) => {
                        add_type_alias_impl(&trigger, acc, ctx, type_item)
                    }
                    hir::AssocItem::Const(const_item) => {
                        add_const_impl(&trigger, acc, ctx, const_item)
                    }
                }),

            SyntaxKind::FN => {
                for missing_fn in get_missing_assoc_items(&ctx.sema, &impl_def)
                    .into_iter()
                    .filter_map(|item| match item {
                        hir::AssocItem::Function(fn_item) => Some(fn_item),
                        _ => None,
                    })
                {
                    add_function_impl(&trigger, acc, ctx, missing_fn);
                }
            }

            SyntaxKind::TYPE_ALIAS => {
                for missing_fn in get_missing_assoc_items(&ctx.sema, &impl_def)
                    .into_iter()
                    .filter_map(|item| match item {
                        hir::AssocItem::TypeAlias(type_item) => Some(type_item),
                        _ => None,
                    })
                {
                    add_type_alias_impl(&trigger, acc, ctx, missing_fn);
                }
            }

            SyntaxKind::CONST => {
                for missing_fn in get_missing_assoc_items(&ctx.sema, &impl_def)
                    .into_iter()
                    .filter_map(|item| match item {
                        hir::AssocItem::Const(const_item) => Some(const_item),
                        _ => None,
                    })
                {
                    add_const_impl(&trigger, acc, ctx, missing_fn);
                }
            }

            _ => {}
        }
    }
}

fn completion_match(ctx: &CompletionContext) -> Option<(SyntaxNode, Impl)> {
    let (trigger, impl_def_offset) = ctx.token.ancestors().find_map(|p| match p.kind() {
        SyntaxKind::FN | SyntaxKind::TYPE_ALIAS | SyntaxKind::CONST | SyntaxKind::BLOCK_EXPR => {
            Some((p, 2))
        }
        SyntaxKind::NAME_REF => Some((p, 5)),
        _ => None,
    })?;
    let impl_def = (0..impl_def_offset - 1)
        .try_fold(trigger.parent()?, |t, _| t.parent())
        .and_then(ast::Impl::cast)?;
    Some((trigger, impl_def))
}

fn add_function_impl(
    fn_def_node: &SyntaxNode,
    acc: &mut Completions,
    ctx: &CompletionContext,
    func: hir::Function,
) {
    let fn_name = func.name(ctx.db).to_string();

    let label = if func.params(ctx.db).is_empty() {
        format!("fn {}()", fn_name)
    } else {
        format!("fn {}(..)", fn_name)
    };

    let builder = CompletionItem::new(CompletionKind::Magic, ctx.source_range(), label)
        .lookup_by(fn_name)
        .set_documentation(func.docs(ctx.db));

    let completion_kind = if func.self_param(ctx.db).is_some() {
        CompletionItemKind::Method
    } else {
        CompletionItemKind::Function
    };
    let range = TextRange::new(fn_def_node.text_range().start(), ctx.source_range().end());

    let function_decl = function_declaration(&func.source(ctx.db).value);
    match ctx.config.snippet_cap {
        Some(cap) => {
            let snippet = format!("{} {{\n    $0\n}}", function_decl);
            builder.snippet_edit(cap, TextEdit::replace(range, snippet))
        }
        None => {
            let header = format!("{} {{", function_decl);
            builder.text_edit(TextEdit::replace(range, header))
        }
    }
    .kind(completion_kind)
    .add_to(acc);
}

fn add_type_alias_impl(
    type_def_node: &SyntaxNode,
    acc: &mut Completions,
    ctx: &CompletionContext,
    type_alias: hir::TypeAlias,
) {
    let alias_name = type_alias.name(ctx.db).to_string();

    let snippet = format!("type {} = ", alias_name);

    let range = TextRange::new(type_def_node.text_range().start(), ctx.source_range().end());

    CompletionItem::new(CompletionKind::Magic, ctx.source_range(), snippet.clone())
        .text_edit(TextEdit::replace(range, snippet))
        .lookup_by(alias_name)
        .kind(CompletionItemKind::TypeAlias)
        .set_documentation(type_alias.docs(ctx.db))
        .add_to(acc);
}

fn add_const_impl(
    const_def_node: &SyntaxNode,
    acc: &mut Completions,
    ctx: &CompletionContext,
    const_: hir::Const,
) {
    let const_name = const_.name(ctx.db).map(|n| n.to_string());

    if let Some(const_name) = const_name {
        let snippet = make_const_compl_syntax(&const_.source(ctx.db).value);

        let range = TextRange::new(const_def_node.text_range().start(), ctx.source_range().end());

        CompletionItem::new(CompletionKind::Magic, ctx.source_range(), snippet.clone())
            .text_edit(TextEdit::replace(range, snippet))
            .lookup_by(const_name)
            .kind(CompletionItemKind::Const)
            .set_documentation(const_.docs(ctx.db))
            .add_to(acc);
    }
}

fn make_const_compl_syntax(const_: &ast::Const) -> String {
    let const_ = edit::remove_attrs_and_docs(const_);

    let const_start = const_.syntax().text_range().start();
    let const_end = const_.syntax().text_range().end();

    let start =
        const_.syntax().first_child_or_token().map_or(const_start, |f| f.text_range().start());

    let end = const_
        .syntax()
        .children_with_tokens()
        .find(|s| s.kind() == T![;] || s.kind() == T![=])
        .map_or(const_end, |f| f.text_range().start());

    let len = end - start;
    let range = TextRange::new(0.into(), len);

    let syntax = const_.syntax().text().slice(range).to_string();

    format!("{} = ", syntax.trim_end())
}

#[cfg(test)]
mod tests {
    use expect_test::{expect, Expect};

    use crate::completion::{
        test_utils::{check_edit, completion_list},
        CompletionKind,
    };

    fn check(ra_fixture: &str, expect: Expect) {
        let actual = completion_list(ra_fixture, CompletionKind::Magic);
        expect.assert_eq(&actual)
    }

    #[test]
    fn name_ref_function_type_const() {
        check(
            r#"
trait Test {
    type TestType;
    const TEST_CONST: u16;
    fn test();
}
struct T;

impl Test for T {
    t<|>
}
"#,
            expect![["
ct const TEST_CONST: u16 = \n\
fn fn test()
ta type TestType = \n\
            "]],
        );
    }

    #[test]
    fn no_nested_fn_completions() {
        check(
            r"
trait Test {
    fn test();
    fn test2();
}
struct T;

impl Test for T {
    fn test() {
        t<|>
    }
}
",
            expect![[""]],
        );
    }

    #[test]
    fn name_ref_single_function() {
        check_edit(
            "test",
            r#"
trait Test {
    fn test();
}
struct T;

impl Test for T {
    t<|>
}
"#,
            r#"
trait Test {
    fn test();
}
struct T;

impl Test for T {
    fn test() {
    $0
}
}
"#,
        );
    }

    #[test]
    fn single_function() {
        check_edit(
            "test",
            r#"
trait Test {
    fn test();
}
struct T;

impl Test for T {
    fn t<|>
}
"#,
            r#"
trait Test {
    fn test();
}
struct T;

impl Test for T {
    fn test() {
    $0
}
}
"#,
        );
    }

    #[test]
    fn hide_implemented_fn() {
        check(
            r#"
trait Test {
    fn foo();
    fn foo_bar();
}
struct T;

impl Test for T {
    fn foo() {}
    fn f<|>
}
"#,
            expect![[r#"
                fn fn foo_bar()
            "#]],
        );
    }

    #[test]
    fn generic_fn() {
        check_edit(
            "foo",
            r#"
trait Test {
    fn foo<T>();
}
struct T;

impl Test for T {
    fn f<|>
}
"#,
            r#"
trait Test {
    fn foo<T>();
}
struct T;

impl Test for T {
    fn foo<T>() {
    $0
}
}
"#,
        );
        check_edit(
            "foo",
            r#"
trait Test {
    fn foo<T>() where T: Into<String>;
}
struct T;

impl Test for T {
    fn f<|>
}
"#,
            r#"
trait Test {
    fn foo<T>() where T: Into<String>;
}
struct T;

impl Test for T {
    fn foo<T>()
where T: Into<String> {
    $0
}
}
"#,
        );
    }

    #[test]
    fn associated_type() {
        check_edit(
            "SomeType",
            r#"
trait Test {
    type SomeType;
}

impl Test for () {
    type S<|>
}
"#,
            "
trait Test {
    type SomeType;
}

impl Test for () {
    type SomeType = \n\
}
",
        );
    }

    #[test]
    fn associated_const() {
        check_edit(
            "SOME_CONST",
            r#"
trait Test {
    const SOME_CONST: u16;
}

impl Test for () {
    const S<|>
}
"#,
            "
trait Test {
    const SOME_CONST: u16;
}

impl Test for () {
    const SOME_CONST: u16 = \n\
}
",
        );

        check_edit(
            "SOME_CONST",
            r#"
trait Test {
    const SOME_CONST: u16 = 92;
}

impl Test for () {
    const S<|>
}
"#,
            "
trait Test {
    const SOME_CONST: u16 = 92;
}

impl Test for () {
    const SOME_CONST: u16 = \n\
}
",
        );
    }
}
<commit_msg>Complete trait impl immediately after type/const/fn<commit_after>//! Completion for associated items in a trait implementation.
//!
//! This module adds the completion items related to implementing associated
//! items within a `impl Trait for Struct` block. The current context node
//! must be within either a `FN`, `TYPE_ALIAS`, or `CONST` node
//! and an direct child of an `IMPL`.
//!
//! # Examples
//!
//! Considering the following trait `impl`:
//!
//! ```ignore
//! trait SomeTrait {
//!     fn foo();
//! }
//!
//! impl SomeTrait for () {
//!     fn f<|>
//! }
//! ```
//!
//! may result in the completion of the following method:
//!
//! ```ignore
//! # trait SomeTrait {
//! #    fn foo();
//! # }
//!
//! impl SomeTrait for () {
//!     fn foo() {}<|>
//! }
//! ```

use assists::utils::get_missing_assoc_items;
use hir::{self, HasAttrs, HasSource};
use syntax::{
    ast::{self, edit, Impl},
    AstNode, SyntaxKind, SyntaxNode, TextRange, T,
};
use text_edit::TextEdit;

use crate::{
    completion::{
        CompletionContext, CompletionItem, CompletionItemKind, CompletionKind, Completions,
    },
    display::function_declaration,
};

#[derive(Debug, PartialEq, Eq)]
enum ImplCompletionKind {
    All,
    Fn,
    TypeAlias,
    Const,
}

pub(crate) fn complete_trait_impl(acc: &mut Completions, ctx: &CompletionContext) {
    if let Some((kind, trigger, impl_def)) = completion_match(ctx) {
        get_missing_assoc_items(&ctx.sema, &impl_def).into_iter().for_each(|item| match item {
            hir::AssocItem::Function(fn_item)
                if kind == ImplCompletionKind::All || kind == ImplCompletionKind::Fn =>
            {
                add_function_impl(&trigger, acc, ctx, fn_item)
            }
            hir::AssocItem::TypeAlias(type_item)
                if kind == ImplCompletionKind::All || kind == ImplCompletionKind::TypeAlias =>
            {
                add_type_alias_impl(&trigger, acc, ctx, type_item)
            }
            hir::AssocItem::Const(const_item)
                if kind == ImplCompletionKind::All || kind == ImplCompletionKind::Const =>
            {
                add_const_impl(&trigger, acc, ctx, const_item)
            }
            _ => {}
        });
    }
}

fn completion_match(ctx: &CompletionContext) -> Option<(ImplCompletionKind, SyntaxNode, Impl)> {
    let mut token = ctx.token.clone();
    // For keywork without name like `impl .. { fn <|> }`, the current position is inside
    // the whitespace token, which is outside `FN` syntax node.
    // We need to follow the previous token in this case.
    if token.kind() == SyntaxKind::WHITESPACE {
        token = token.prev_token()?;
    }

    let (kind, trigger, impl_def_offset) = token.ancestors().find_map(|p| match p.kind() {
        // `const` can be a modifier of an item, so the `const` token may be inside another item syntax node.
        // Eg. `impl .. { const <|> fn bar() .. }`
        SyntaxKind::FN | SyntaxKind::TYPE_ALIAS | SyntaxKind::CONST
            if token.kind() == SyntaxKind::CONST_KW =>
        {
            Some((ImplCompletionKind::Const, p, 2))
        }
        SyntaxKind::FN => Some((ImplCompletionKind::Fn, p, 2)),
        SyntaxKind::TYPE_ALIAS => Some((ImplCompletionKind::TypeAlias, p, 2)),
        SyntaxKind::CONST => Some((ImplCompletionKind::Const, p, 2)),
        // `impl .. { const <|> }` is parsed as:
        // IMPL
        //   ASSOC_ITEM_LIST
        //     ERROR
        //       CONST_KW <- token
        //     WHITESPACE <- ctx.token
        SyntaxKind::ERROR
            if p.first_token().map_or(false, |t| t.kind() == SyntaxKind::CONST_KW) =>
        {
            Some((ImplCompletionKind::Const, p, 2))
        }
        SyntaxKind::NAME_REF => Some((ImplCompletionKind::All, p, 5)),
        _ => None,
    })?;

    let impl_def = (0..impl_def_offset - 1)
        .try_fold(trigger.parent()?, |t, _| t.parent())
        .and_then(ast::Impl::cast)?;
    Some((kind, trigger, impl_def))
}

fn add_function_impl(
    fn_def_node: &SyntaxNode,
    acc: &mut Completions,
    ctx: &CompletionContext,
    func: hir::Function,
) {
    let fn_name = func.name(ctx.db).to_string();

    let label = if func.params(ctx.db).is_empty() {
        format!("fn {}()", fn_name)
    } else {
        format!("fn {}(..)", fn_name)
    };

    let builder = CompletionItem::new(CompletionKind::Magic, ctx.source_range(), label)
        .lookup_by(fn_name)
        .set_documentation(func.docs(ctx.db));

    let completion_kind = if func.self_param(ctx.db).is_some() {
        CompletionItemKind::Method
    } else {
        CompletionItemKind::Function
    };
    let range = TextRange::new(fn_def_node.text_range().start(), ctx.source_range().end());

    let function_decl = function_declaration(&func.source(ctx.db).value);
    match ctx.config.snippet_cap {
        Some(cap) => {
            let snippet = format!("{} {{\n    $0\n}}", function_decl);
            builder.snippet_edit(cap, TextEdit::replace(range, snippet))
        }
        None => {
            let header = format!("{} {{", function_decl);
            builder.text_edit(TextEdit::replace(range, header))
        }
    }
    .kind(completion_kind)
    .add_to(acc);
}

fn add_type_alias_impl(
    type_def_node: &SyntaxNode,
    acc: &mut Completions,
    ctx: &CompletionContext,
    type_alias: hir::TypeAlias,
) {
    let alias_name = type_alias.name(ctx.db).to_string();

    let snippet = format!("type {} = ", alias_name);

    let range = TextRange::new(type_def_node.text_range().start(), ctx.source_range().end());

    CompletionItem::new(CompletionKind::Magic, ctx.source_range(), snippet.clone())
        .text_edit(TextEdit::replace(range, snippet))
        .lookup_by(alias_name)
        .kind(CompletionItemKind::TypeAlias)
        .set_documentation(type_alias.docs(ctx.db))
        .add_to(acc);
}

fn add_const_impl(
    const_def_node: &SyntaxNode,
    acc: &mut Completions,
    ctx: &CompletionContext,
    const_: hir::Const,
) {
    let const_name = const_.name(ctx.db).map(|n| n.to_string());

    if let Some(const_name) = const_name {
        let snippet = make_const_compl_syntax(&const_.source(ctx.db).value);

        let range = TextRange::new(const_def_node.text_range().start(), ctx.source_range().end());

        CompletionItem::new(CompletionKind::Magic, ctx.source_range(), snippet.clone())
            .text_edit(TextEdit::replace(range, snippet))
            .lookup_by(const_name)
            .kind(CompletionItemKind::Const)
            .set_documentation(const_.docs(ctx.db))
            .add_to(acc);
    }
}

fn make_const_compl_syntax(const_: &ast::Const) -> String {
    let const_ = edit::remove_attrs_and_docs(const_);

    let const_start = const_.syntax().text_range().start();
    let const_end = const_.syntax().text_range().end();

    let start =
        const_.syntax().first_child_or_token().map_or(const_start, |f| f.text_range().start());

    let end = const_
        .syntax()
        .children_with_tokens()
        .find(|s| s.kind() == T![;] || s.kind() == T![=])
        .map_or(const_end, |f| f.text_range().start());

    let len = end - start;
    let range = TextRange::new(0.into(), len);

    let syntax = const_.syntax().text().slice(range).to_string();

    format!("{} = ", syntax.trim_end())
}

#[cfg(test)]
mod tests {
    use expect_test::{expect, Expect};

    use crate::completion::{
        test_utils::{check_edit, completion_list},
        CompletionKind,
    };

    fn check(ra_fixture: &str, expect: Expect) {
        let actual = completion_list(ra_fixture, CompletionKind::Magic);
        expect.assert_eq(&actual)
    }

    #[test]
    fn name_ref_function_type_const() {
        check(
            r#"
trait Test {
    type TestType;
    const TEST_CONST: u16;
    fn test();
}
struct T;

impl Test for T {
    t<|>
}
"#,
            expect![["
ct const TEST_CONST: u16 = \n\
fn fn test()
ta type TestType = \n\
            "]],
        );
    }

    #[test]
    fn no_nested_fn_completions() {
        check(
            r"
trait Test {
    fn test();
    fn test2();
}
struct T;

impl Test for T {
    fn test() {
        t<|>
    }
}
",
            expect![[""]],
        );
    }

    #[test]
    fn name_ref_single_function() {
        check_edit(
            "test",
            r#"
trait Test {
    fn test();
}
struct T;

impl Test for T {
    t<|>
}
"#,
            r#"
trait Test {
    fn test();
}
struct T;

impl Test for T {
    fn test() {
    $0
}
}
"#,
        );
    }

    #[test]
    fn single_function() {
        check_edit(
            "test",
            r#"
trait Test {
    fn test();
}
struct T;

impl Test for T {
    fn t<|>
}
"#,
            r#"
trait Test {
    fn test();
}
struct T;

impl Test for T {
    fn test() {
    $0
}
}
"#,
        );
    }

    #[test]
    fn hide_implemented_fn() {
        check(
            r#"
trait Test {
    fn foo();
    fn foo_bar();
}
struct T;

impl Test for T {
    fn foo() {}
    fn f<|>
}
"#,
            expect![[r#"
                fn fn foo_bar()
            "#]],
        );
    }

    #[test]
    fn generic_fn() {
        check_edit(
            "foo",
            r#"
trait Test {
    fn foo<T>();
}
struct T;

impl Test for T {
    fn f<|>
}
"#,
            r#"
trait Test {
    fn foo<T>();
}
struct T;

impl Test for T {
    fn foo<T>() {
    $0
}
}
"#,
        );
        check_edit(
            "foo",
            r#"
trait Test {
    fn foo<T>() where T: Into<String>;
}
struct T;

impl Test for T {
    fn f<|>
}
"#,
            r#"
trait Test {
    fn foo<T>() where T: Into<String>;
}
struct T;

impl Test for T {
    fn foo<T>()
where T: Into<String> {
    $0
}
}
"#,
        );
    }

    #[test]
    fn associated_type() {
        check_edit(
            "SomeType",
            r#"
trait Test {
    type SomeType;
}

impl Test for () {
    type S<|>
}
"#,
            "
trait Test {
    type SomeType;
}

impl Test for () {
    type SomeType = \n\
}
",
        );
    }

    #[test]
    fn associated_const() {
        check_edit(
            "SOME_CONST",
            r#"
trait Test {
    const SOME_CONST: u16;
}

impl Test for () {
    const S<|>
}
"#,
            "
trait Test {
    const SOME_CONST: u16;
}

impl Test for () {
    const SOME_CONST: u16 = \n\
}
",
        );

        check_edit(
            "SOME_CONST",
            r#"
trait Test {
    const SOME_CONST: u16 = 92;
}

impl Test for () {
    const S<|>
}
"#,
            "
trait Test {
    const SOME_CONST: u16 = 92;
}

impl Test for () {
    const SOME_CONST: u16 = \n\
}
",
        );
    }

    #[test]
    fn complete_without_name() {
        let test = |completion: &str, hint: &str, completed: &str, next_sibling: &str| {
            println!(
                "completion='{}', hint='{}', next_sibling='{}'",
                completion, hint, next_sibling
            );

            check_edit(
                completion,
                &format!(
                    r#"
trait Test {{
    type Foo;
    const CONST: u16;
    fn bar();
}}
struct T;

impl Test for T {{
    {}
    {}
}}
"#,
                    hint, next_sibling
                ),
                &format!(
                    r#"
trait Test {{
    type Foo;
    const CONST: u16;
    fn bar();
}}
struct T;

impl Test for T {{
    {}
    {}
}}
"#,
                    completed, next_sibling
                ),
            )
        };

        // Enumerate some possible next siblings.
        for next_sibling in &[
            "",
            "fn other_fn() {}", // `const <|> fn` -> `const fn`
            "type OtherType = i32;",
            "const OTHER_CONST: i32 = 0;",
            "async fn other_fn() {}",
            "unsafe fn other_fn() {}",
            "default fn other_fn() {}",
            "default type OtherType = i32;",
            "default const OTHER_CONST: i32 = 0;",
        ] {
            test("bar", "fn <|>", "fn bar() {\n    $0\n}", next_sibling);
            test("Foo", "type <|>", "type Foo = ", next_sibling);
            test("CONST", "const <|>", "const CONST: u16 = ", next_sibling);
        }
    }
}
<|endoftext|>