<commit_before>use crate::cell::UnsafeCell;
use crate::mem;
use crate::sync::atomic::{AtomicU32, Ordering};
use crate::sys::cloudabi::abi;
use crate::sys::rwlock::{self, RWLock};

extern "C" {
    #[thread_local]
    static __pthread_thread_id: abi::tid;
}

// Implement Mutex using an RWLock. This doesn't introduce any
// performance overhead in this environment, as the operations would be
// implemented identically.
pub struct Mutex(RWLock);

pub unsafe fn raw(m: &Mutex) -> *mut AtomicU32 {
    rwlock::raw(&m.0)
}

impl Mutex {
    pub const fn new() -> Mutex {
        Mutex(RWLock::new())
    }

    pub unsafe fn init(&mut self) {
        // This function should normally reinitialize the mutex after
        // moving it to a different memory address. This implementation
        // does not require adjustments after moving.
    }

    pub unsafe fn try_lock(&self) -> bool {
        self.0.try_write()
    }

    pub unsafe fn lock(&self) {
        self.0.write()
    }

    pub unsafe fn unlock(&self) {
        self.0.write_unlock()
    }

    pub unsafe fn destroy(&self) {
        self.0.destroy()
    }
}

pub struct ReentrantMutex {
    lock: UnsafeCell<AtomicU32>,
    recursion: UnsafeCell<u32>,
}

impl ReentrantMutex {
    pub unsafe fn uninitialized() -> ReentrantMutex {
        mem::uninitialized()
    }

    pub unsafe fn init(&mut self) {
        self.lock = UnsafeCell::new(AtomicU32::new(abi::LOCK_UNLOCKED.0));
        self.recursion = UnsafeCell::new(0);
    }

    pub unsafe fn try_lock(&self) -> bool {
        // Attempt to acquire the lock.
        let lock = self.lock.get();
        let recursion = self.recursion.get();
        if let Err(old) = (*lock).compare_exchange(
            abi::LOCK_UNLOCKED.0,
            __pthread_thread_id.0 | abi::LOCK_WRLOCKED.0,
            Ordering::Acquire,
            Ordering::Relaxed,
        ) {
            // If we fail to acquire the lock, it may be the case
            // that we've already acquired it and may need to recurse.
            if old & !abi::LOCK_KERNEL_MANAGED.0 == __pthread_thread_id.0 | abi::LOCK_WRLOCKED.0 {
                *recursion += 1;
                true
            } else {
                false
            }
        } else {
            // Success.
            assert_eq!(*recursion, 0, "Mutex has invalid recursion count");
            true
        }
    }

    pub unsafe fn lock(&self) {
        if !self.try_lock() {
            // Call into the kernel to acquire a write lock.
            let lock = self.lock.get();
            let subscription = abi::subscription {
                type_: abi::eventtype::LOCK_WRLOCK,
                union: abi::subscription_union {
                    lock: abi::subscription_lock {
                        lock: lock as *mut abi::lock,
                        lock_scope: abi::scope::PRIVATE,
                    },
                },
                ..mem::zeroed()
            };
            let mut event: abi::event = mem::uninitialized();
            let mut nevents: usize = mem::uninitialized();
            let ret = abi::poll(&subscription, &mut event, 1, &mut nevents);
            assert_eq!(ret, abi::errno::SUCCESS, "Failed to acquire mutex");
            assert_eq!(event.error, abi::errno::SUCCESS, "Failed to acquire mutex");
        }
    }

    pub unsafe fn unlock(&self) {
        let lock = self.lock.get();
        let recursion = self.recursion.get();
        assert_eq!(
            (*lock).load(Ordering::Relaxed) & !abi::LOCK_KERNEL_MANAGED.0,
            __pthread_thread_id.0 | abi::LOCK_WRLOCKED.0,
            "This mutex is locked by a different thread"
        );

        if *recursion > 0 {
            *recursion -= 1;
        } else if !(*lock)
            .compare_exchange(
                __pthread_thread_id.0 | abi::LOCK_WRLOCKED.0,
                abi::LOCK_UNLOCKED.0,
                Ordering::Release,
                Ordering::Relaxed,
            )
            .is_ok()
        {
            // Lock is managed by kernelspace. Call into the kernel
            // to unblock waiting threads.
            let ret = abi::lock_unlock(lock as *mut abi::lock, abi::scope::PRIVATE);
            assert_eq!(ret, abi::errno::SUCCESS, "Failed to unlock a mutex");
        }
    }

    pub unsafe fn destroy(&self) {
        let lock = self.lock.get();
        let recursion = self.recursion.get();
        assert_eq!(
            (*lock).load(Ordering::Relaxed),
            abi::LOCK_UNLOCKED.0,
            "Attempted to destroy locked mutex"
        );
        assert_eq!(*recursion, 0, "Recursion counter invalid");
    }
}
<commit_msg>Modify CloudABI ReentrantMutex to use MaybeUninit<commit_after>use crate::cell::UnsafeCell;
use crate::mem;
use crate::mem::MaybeUninit;
use crate::sync::atomic::{AtomicU32, Ordering};
use crate::sys::cloudabi::abi;
use crate::sys::rwlock::{self, RWLock};

extern "C" {
    #[thread_local]
    static __pthread_thread_id: abi::tid;
}

// Implement Mutex using an RWLock. This doesn't introduce any
// performance overhead in this environment, as the operations would be
// implemented identically.
pub struct Mutex(RWLock);

pub unsafe fn raw(m: &Mutex) -> *mut AtomicU32 {
    rwlock::raw(&m.0)
}

impl Mutex {
    pub const fn new() -> Mutex {
        Mutex(RWLock::new())
    }

    pub unsafe fn init(&mut self) {
        // This function should normally reinitialize the mutex after
        // moving it to a different memory address. This implementation
        // does not require adjustments after moving.
    }

    pub unsafe fn try_lock(&self) -> bool {
        self.0.try_write()
    }

    pub unsafe fn lock(&self) {
        self.0.write()
    }

    pub unsafe fn unlock(&self) {
        self.0.write_unlock()
    }

    pub unsafe fn destroy(&self) {
        self.0.destroy()
    }
}

pub struct ReentrantMutex {
    lock: UnsafeCell<MaybeUninit<AtomicU32>>,
    recursion: UnsafeCell<MaybeUninit<u32>>,
}

impl ReentrantMutex {
    pub unsafe fn uninitialized() -> ReentrantMutex {
        ReentrantMutex {
            lock: UnsafeCell::new(MaybeUninit::uninit()),
            recursion: UnsafeCell::new(MaybeUninit::uninit())
        }
    }

    pub unsafe fn init(&mut self) {
        self.lock = UnsafeCell::new(MaybeUninit::new(AtomicU32::new(abi::LOCK_UNLOCKED.0)));
        self.recursion = UnsafeCell::new(MaybeUninit::new(0));
    }

    pub unsafe fn try_lock(&self) -> bool {
        // Attempt to acquire the lock.
        let lock = self.lock.get();
        let recursion = self.recursion.get();
        if let Err(old) = (*(*lock).as_mut_ptr()).compare_exchange(
            abi::LOCK_UNLOCKED.0,
            __pthread_thread_id.0 | abi::LOCK_WRLOCKED.0,
            Ordering::Acquire,
            Ordering::Relaxed,
        ) {
            // If we fail to acquire the lock, it may be the case
            // that we've already acquired it and may need to recurse.
            if old & !abi::LOCK_KERNEL_MANAGED.0 == __pthread_thread_id.0 | abi::LOCK_WRLOCKED.0 {
                *(*recursion).as_mut_ptr() += 1;
                true
            } else {
                false
            }
        } else {
            // Success.
            assert_eq!(*(*recursion).as_mut_ptr(), 0, "Mutex has invalid recursion count");
            true
        }
    }

    pub unsafe fn lock(&self) {
        if !self.try_lock() {
            // Call into the kernel to acquire a write lock.
            let lock = self.lock.get();
            let subscription = abi::subscription {
                type_: abi::eventtype::LOCK_WRLOCK,
                union: abi::subscription_union {
                    lock: abi::subscription_lock {
                        lock: lock as *mut abi::lock,
                        lock_scope: abi::scope::PRIVATE,
                    },
                },
                ..mem::zeroed()
            };
            let mut event: abi::event = mem::uninitialized();
            let mut nevents: usize = mem::uninitialized();
            let ret = abi::poll(&subscription, &mut event, 1, &mut nevents);
            assert_eq!(ret, abi::errno::SUCCESS, "Failed to acquire mutex");
            assert_eq!(event.error, abi::errno::SUCCESS, "Failed to acquire mutex");
        }
    }

    pub unsafe fn unlock(&self) {
        let lock = self.lock.get();
        let recursion = self.recursion.get();
        assert_eq!(
            (*(*lock).as_mut_ptr()).load(Ordering::Relaxed) & !abi::LOCK_KERNEL_MANAGED.0,
            __pthread_thread_id.0 | abi::LOCK_WRLOCKED.0,
            "This mutex is locked by a different thread"
        );

        if *(*recursion).as_mut_ptr() > 0 {
            *(*recursion).as_mut_ptr() -= 1;
        } else if !(*(*lock).as_mut_ptr())
            .compare_exchange(
                __pthread_thread_id.0 | abi::LOCK_WRLOCKED.0,
                abi::LOCK_UNLOCKED.0,
                Ordering::Release,
                Ordering::Relaxed,
            )
            .is_ok()
        {
            // Lock is managed by kernelspace. Call into the kernel
            // to unblock waiting threads.
            let ret = abi::lock_unlock(lock as *mut abi::lock, abi::scope::PRIVATE);
            assert_eq!(ret, abi::errno::SUCCESS, "Failed to unlock a mutex");
        }
    }

    pub unsafe fn destroy(&self) {
        let lock = self.lock.get();
        let recursion = self.recursion.get();
        assert_eq!(
            (*(*lock).as_mut_ptr()).load(Ordering::Relaxed),
            abi::LOCK_UNLOCKED.0,
            "Attempted to destroy locked mutex"
        );
        assert_eq!(*(*recursion).as_mut_ptr(), 0, "Recursion counter invalid");
    }
}
<|endoftext|>


<commit_before>use crate::{
    builders::CrateBuilder,
    util::{MockAnonymousUser, RequestHelper, TestApp, TestDatabase},
};
use http::StatusCode;
use std::time::Duration;

#[test]
fn download_crate_with_broken_networking_primary_database() {
    let (app, anon, _, owner) = TestApp::init()
        .with_database(TestDatabase::SlowRealPool { replica: false })
        .with_token();
    app.db(|conn| {
        CrateBuilder::new("crate_name", owner.as_model().user_id)
            .version("1.0.0")
            .expect_build(conn)
    });

    // When the database connection is healthy downloads are redirected with the proper
    // capitalization, and missing crates or versions return a 404.

    assert_checked_redirects(&anon);

    // After networking breaks, preventing new database connections, the download endpoint should
    // do an unconditional redirect to the CDN, without checking whether the crate exists or what
    // the exact capitalization of crate name is.

    app.primary_db_chaosproxy().break_networking();
    assert_unconditional_redirects(&anon);

    // After restoring the network and waiting for the database pool to get healthy again redirects
    // should be checked again.

    app.primary_db_chaosproxy().restore_networking();
    app.as_inner()
        .primary_database
        .wait_until_healthy(Duration::from_millis(500))
        .expect("the database did not return healthy");

    assert_checked_redirects(&anon);
}

fn assert_checked_redirects(anon: &MockAnonymousUser) {
    anon.get::<()>("/api/v1/crates/crate_name/1.0.0/download")
        .assert_redirect_ends_with("/crate_name/crate_name-1.0.0.crate");

    anon.get::<()>("/api/v1/crates/Crate-Name/1.0.0/download")
        .assert_redirect_ends_with("/crate_name/crate_name-1.0.0.crate");

    anon.get::<()>("/api/v1/crates/crate_name/2.0.0/download")
        .assert_not_found();

    anon.get::<()>("/api/v1/crates/awesome-project/1.0.0/download")
        .assert_not_found();
}

fn assert_unconditional_redirects(anon: &MockAnonymousUser) {
    anon.get::<()>("/api/v1/crates/crate_name/1.0.0/download")
        .assert_redirect_ends_with("/crate_name/crate_name-1.0.0.crate");

    anon.get::<()>("/api/v1/crates/Crate-Name/1.0.0/download")
        .assert_redirect_ends_with("/Crate-Name/Crate-Name-1.0.0.crate");

    anon.get::<()>("/api/v1/crates/crate_name/2.0.0/download")
        .assert_redirect_ends_with("/crate_name/crate_name-2.0.0.crate");

    anon.get::<()>("/api/v1/crates/awesome-project/1.0.0/download")
        .assert_redirect_ends_with("/awesome-project/awesome-project-1.0.0.crate");
}

#[test]
fn http_error_with_unhealthy_database() {
    let (app, anon) = TestApp::init()
        .with_database(TestDatabase::SlowRealPool { replica: false })
        .empty();

    let response = anon.get::<()>("/api/v1/summary");
    assert_eq!(response.status(), StatusCode::OK);

    app.primary_db_chaosproxy().break_networking();

    let response = anon.get::<()>("/api/v1/summary");
    assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);

    app.primary_db_chaosproxy().restore_networking();
    app.as_inner()
        .primary_database
        .wait_until_healthy(Duration::from_millis(500))
        .expect("the database did not return healthy");

    let response = anon.get::<()>("/api/v1/summary");
    assert_eq!(response.status(), StatusCode::OK);
}

#[test]
fn fallback_to_replica_returns_user_info() {
    const URL: &str = "/api/v1/users/foo";

    let (app, _, owner) = TestApp::init()
        .with_database(TestDatabase::SlowRealPool { replica: true })
        .with_user();
    app.db_new_user("foo");
    app.primary_db_chaosproxy().break_networking();

    // When the primary database is down, requests are forwarded to the replica database
    let response = owner.get::<()>(URL);
    assert_eq!(response.status(), 200);

    // restore primary database connection
    app.primary_db_chaosproxy().restore_networking();
    app.as_inner()
        .primary_database
        .wait_until_healthy(Duration::from_millis(500))
        .expect("the database did not return healthy");
}

#[test]
fn restored_replica_returns_user_info() {
    const URL: &str = "/api/v1/users/foo";

    let (app, _, owner) = TestApp::init()
        .with_database(TestDatabase::SlowRealPool { replica: true })
        .with_user();
    app.db_new_user("foo");
    app.primary_db_chaosproxy().break_networking();
    app.replica_db_chaosproxy().break_networking();

    // When both primary and replica database are down, the request returns an error
    let response = owner.get::<()>(URL);
    assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);

    // Once the replica database is restored, it should serve as a fallback again
    app.replica_db_chaosproxy().restore_networking();
    app.as_inner()
        .read_only_replica_database
        .as_ref()
        .expect("no replica database configured")
        .wait_until_healthy(Duration::from_millis(500))
        .expect("the database did not return healthy");

    let response = owner.get::<()>(URL);
    assert_eq!(response.status(), StatusCode::OK);

    // restore connection
    app.primary_db_chaosproxy().restore_networking();
    app.as_inner()
        .primary_database
        .wait_until_healthy(Duration::from_millis(500))
        .expect("the database did not return healthy");
}
<commit_msg>Add test to check recovery of primary connection<commit_after>use crate::{
    builders::CrateBuilder,
    util::{MockAnonymousUser, RequestHelper, TestApp, TestDatabase},
};
use http::StatusCode;
use std::time::Duration;

#[test]
fn download_crate_with_broken_networking_primary_database() {
    let (app, anon, _, owner) = TestApp::init()
        .with_database(TestDatabase::SlowRealPool { replica: false })
        .with_token();
    app.db(|conn| {
        CrateBuilder::new("crate_name", owner.as_model().user_id)
            .version("1.0.0")
            .expect_build(conn)
    });

    // When the database connection is healthy downloads are redirected with the proper
    // capitalization, and missing crates or versions return a 404.

    assert_checked_redirects(&anon);

    // After networking breaks, preventing new database connections, the download endpoint should
    // do an unconditional redirect to the CDN, without checking whether the crate exists or what
    // the exact capitalization of crate name is.

    app.primary_db_chaosproxy().break_networking();
    assert_unconditional_redirects(&anon);

    // After restoring the network and waiting for the database pool to get healthy again redirects
    // should be checked again.

    app.primary_db_chaosproxy().restore_networking();
    app.as_inner()
        .primary_database
        .wait_until_healthy(Duration::from_millis(500))
        .expect("the database did not return healthy");

    assert_checked_redirects(&anon);
}

fn assert_checked_redirects(anon: &MockAnonymousUser) {
    anon.get::<()>("/api/v1/crates/crate_name/1.0.0/download")
        .assert_redirect_ends_with("/crate_name/crate_name-1.0.0.crate");

    anon.get::<()>("/api/v1/crates/Crate-Name/1.0.0/download")
        .assert_redirect_ends_with("/crate_name/crate_name-1.0.0.crate");

    anon.get::<()>("/api/v1/crates/crate_name/2.0.0/download")
        .assert_not_found();

    anon.get::<()>("/api/v1/crates/awesome-project/1.0.0/download")
        .assert_not_found();
}

fn assert_unconditional_redirects(anon: &MockAnonymousUser) {
    anon.get::<()>("/api/v1/crates/crate_name/1.0.0/download")
        .assert_redirect_ends_with("/crate_name/crate_name-1.0.0.crate");

    anon.get::<()>("/api/v1/crates/Crate-Name/1.0.0/download")
        .assert_redirect_ends_with("/Crate-Name/Crate-Name-1.0.0.crate");

    anon.get::<()>("/api/v1/crates/crate_name/2.0.0/download")
        .assert_redirect_ends_with("/crate_name/crate_name-2.0.0.crate");

    anon.get::<()>("/api/v1/crates/awesome-project/1.0.0/download")
        .assert_redirect_ends_with("/awesome-project/awesome-project-1.0.0.crate");
}

#[test]
fn http_error_with_unhealthy_database() {
    let (app, anon) = TestApp::init()
        .with_database(TestDatabase::SlowRealPool { replica: false })
        .empty();

    let response = anon.get::<()>("/api/v1/summary");
    assert_eq!(response.status(), StatusCode::OK);

    app.primary_db_chaosproxy().break_networking();

    let response = anon.get::<()>("/api/v1/summary");
    assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);

    app.primary_db_chaosproxy().restore_networking();
    app.as_inner()
        .primary_database
        .wait_until_healthy(Duration::from_millis(500))
        .expect("the database did not return healthy");

    let response = anon.get::<()>("/api/v1/summary");
    assert_eq!(response.status(), StatusCode::OK);
}

#[test]
fn fallback_to_replica_returns_user_info() {
    const URL: &str = "/api/v1/users/foo";

    let (app, _, owner) = TestApp::init()
        .with_database(TestDatabase::SlowRealPool { replica: true })
        .with_user();
    app.db_new_user("foo");
    app.primary_db_chaosproxy().break_networking();

    // When the primary database is down, requests are forwarded to the replica database
    let response = owner.get::<()>(URL);
    assert_eq!(response.status(), 200);

    // restore primary database connection
    app.primary_db_chaosproxy().restore_networking();
    app.as_inner()
        .primary_database
        .wait_until_healthy(Duration::from_millis(500))
        .expect("the database did not return healthy");
}

#[test]
fn restored_replica_returns_user_info() {
    const URL: &str = "/api/v1/users/foo";

    let (app, _, owner) = TestApp::init()
        .with_database(TestDatabase::SlowRealPool { replica: true })
        .with_user();
    app.db_new_user("foo");
    app.primary_db_chaosproxy().break_networking();
    app.replica_db_chaosproxy().break_networking();

    // When both primary and replica database are down, the request returns an error
    let response = owner.get::<()>(URL);
    assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);

    // Once the replica database is restored, it should serve as a fallback again
    app.replica_db_chaosproxy().restore_networking();
    app.as_inner()
        .read_only_replica_database
        .as_ref()
        .expect("no replica database configured")
        .wait_until_healthy(Duration::from_millis(500))
        .expect("the database did not return healthy");

    let response = owner.get::<()>(URL);
    assert_eq!(response.status(), StatusCode::OK);

    // restore connection
    app.primary_db_chaosproxy().restore_networking();
    app.as_inner()
        .primary_database
        .wait_until_healthy(Duration::from_millis(500))
        .expect("the database did not return healthy");
}

#[test]
fn restored_primary_returns_user_info() {
    const URL: &str = "/api/v1/users/foo";

    let (app, _, owner) = TestApp::init().with_database(TestDatabase::SlowRealPool{ replica: true }).with_user();
    app.db_new_user("foo");
    app.primary_db_chaosproxy().break_networking();
    app.replica_db_chaosproxy().break_networking();

    // When both primary and replica database are down, the request returns an error
    let response = owner.get::<()>(URL);
    assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);

    // Once the replica database is restored, it should serve as a fallback again
    app.primary_db_chaosproxy().restore_networking();
    app.as_inner()
        .primary_database
        .wait_until_healthy(Duration::from_millis(500))
        .expect("the database did not return healthy");

    let response = owner.get::<()>(URL);
    assert_eq!(response.status(), StatusCode::OK);
}
<|endoftext|>


<commit_before>use clap::{App, Arg};

pub fn cli_app<'a, 'b>() -> App<'a, 'b> {
    App::new("beacon_node")
        .visible_aliases(&["b", "bn", "beacon"])
        .version(crate_version!())
        .author("Sigma Prime <contact@sigmaprime.io>")
        .about("The primary component which connects to the Ethereum 2.0 P2P network and \
                downloads, verifies and stores blocks. Provides a HTTP API for querying \
                the beacon chain and publishing messages to the network.")
        /*
         * Configuration directory locations.
         */
        .arg(
            Arg::with_name("network-dir")
                .long("network-dir")
                .value_name("DIR")
                .help("Data directory for network keys. Defaults to network/ inside the beacon node \
                       dir.")
                .takes_value(true)
        )
        .arg(
            Arg::with_name("freezer-dir")
                .long("freezer-dir")
                .value_name("DIR")
                .help("Data directory for the freezer database.")
                .takes_value(true)
        )
        /*
         * Network parameters.
         */
        .arg(
            Arg::with_name("zero-ports")
                .long("zero-ports")
                .short("z")
                .help("Sets all listening TCP/UDP ports to 0, allowing the OS to choose some \
                       arbitrary free ports.")
                .takes_value(false),
        )
        .arg(
            Arg::with_name("listen-address")
                .long("listen-address")
                .value_name("ADDRESS")
                .help("The address lighthouse will listen for UDP and TCP connections.")
                .default_value("0.0.0.0")
                .takes_value(true)
        )
        .arg(
            Arg::with_name("port")
                .long("port")
                .value_name("PORT")
                .help("The TCP/UDP port to listen on. The UDP port can be modified by the --discovery-port flag.")
                .default_value("9000")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("discovery-port")
                .long("discovery-port")
                .value_name("PORT")
                .help("The UDP port that discovery will listen on. Defaults to `port`")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("maxpeers")
                .long("maxpeers")
                .help("The maximum number of peers.")
                .default_value("50")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("boot-nodes")
                .long("boot-nodes")
                .allow_hyphen_values(true)
                .value_name("ENR-LIST")
                .help("One or more comma-delimited base64-encoded ENR's to bootstrap the p2p network.")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("enr-udp-port")
                .long("enr-udp-port")
                .value_name("PORT")
                .help("The UDP port of the local ENR. Set this only if you are sure other nodes can connect to your local node on this port.")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("enr-tcp-port")
                .long("enr-tcp-port")
                .value_name("PORT")
                .help("The TCP port of the local ENR. Set this only if you are sure other nodes can connect to your local node on this port.\
                    The --port flag is used if this is not set.")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("enr-address")
                .long("enr-address")
                .value_name("ADDRESS")
                .help("The IP address/ DNS address to broadcast to other peers on how to reach this node. \
                If a DNS address is provided, the enr-address is set to the IP address it resolves to and \
                does not auto-update based on PONG responses in discovery. \
                Set this only if you are sure other nodes can connect to your local node on this address. \
                Discovery will automatically find your external address,if possible.")
                .requires("enr-udp-port")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("enr-match")
                .short("e")
                .long("enr-match")
                .help("Sets the local ENR IP address and port to match those set for lighthouse. \
                Specifically, the IP address will be the value of --listen-address and the UDP port will be --discovery-port.")
        )
        .arg(
            Arg::with_name("disable-enr-auto-update")
                .short("x")
                .long("disable-enr-auto-update")
                .help("Discovery automatically updates the nodes local ENR with an external IP address and port as seen by other peers on the network. \
                This disables this feature, fixing the ENR's IP/PORT to those specified on boot.")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("libp2p-addresses")
                .long("libp2p-addresses")
                .value_name("MULTIADDR")
                .help("One or more comma-delimited multiaddrs to manually connect to a libp2p peer \
                       without an ENR.")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("p2p-priv-key")
                .long("p2p-priv-key")
                .value_name("HEX")
                .help("A secp256k1 secret key, represented as ASCII-encoded hex bytes (with or \
                       without 0x prefix). Default is either loaded from disk or generated \
                       automatically.")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("random-propagation")
                .long("random-propagation")
                .value_name("INTEGER")
                .takes_value(true)
                .help("Specifies (as a percentage) the likelihood of propagating blocks and \
                       attestations. This should only be used for testing networking elements. The \
                       value must like in the range 1-100. Default is 100.")
        )
        /* REST API related arguments */
        .arg(
            Arg::with_name("http")
                .long("http")
                .help("Enable RESTful HTTP API server. Disabled by default.")
                .takes_value(false),
        )
        .arg(
            Arg::with_name("http-address")
                .long("http-address")
                .value_name("ADDRESS")
                .help("Set the listen address for the RESTful HTTP API server.")
                .default_value("127.0.0.1")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("http-port")
                .long("http-port")
                .value_name("PORT")
                .help("Set the listen TCP port for the RESTful HTTP API server.")
                .default_value("5052")
                .takes_value(true),
        )
        /* Websocket related arguments */
        .arg(
            Arg::with_name("ws")
                .long("ws")
                .help("Enable the websocket server. Disabled by default.")
                .takes_value(false),
        )
        .arg(
            Arg::with_name("ws-address")
                .long("ws-address")
                .value_name("ADDRESS")
                .help("Set the listen address for the websocket server.")
                .default_value("127.0.0.1")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("ws-port")
                .long("ws-port")
                .value_name("PORT")
                .help("Set the listen TCP port for the websocket server.")
                .default_value("5053")
                .takes_value(true),
        )

        /*
         * Eth1 Integration
         */
        .arg(
            Arg::with_name("eth1")
                .long("eth1")
                .help("If present the node will connect to an eth1 node. This is required for \
                       block production, you must use this flag if you wish to serve a validator.")
                .takes_value(false),
        )
        .arg(
            Arg::with_name("dummy-eth1")
                .long("dummy-eth1")
                .conflicts_with("eth1")
                .help("If present, uses an eth1 backend that generates static dummy data.\
                      Identical to the method used at the 2019 Canada interop.")
        )
        .arg(
            Arg::with_name("eth1-endpoint")
                .long("eth1-endpoint")
                .value_name("HTTP-ENDPOINT")
                .help("Specifies the server for a web3 connection to the Eth1 chain. Also enables the --eth1 flag.")
                .takes_value(true)
                .default_value("http://127.0.0.1:8545")
        )
        .arg(
            Arg::with_name("slots-per-restore-point")
                .long("slots-per-restore-point")
                .value_name("SLOT_COUNT")
                .help("Specifies how often a freezer DB restore point should be stored. \
                       DO NOT DECREASE AFTER INITIALIZATION. [default: 2048 (mainnet) or 64 (minimal)]")
                .takes_value(true)
        )
        .arg(
            Arg::with_name("block-cache-size")
                .long("block-cache-size")
                .value_name("SIZE")
                .help("Specifies how many blocks the database should cache in memory [default: 5]")
                .takes_value(true)
        )
        .arg(
            Arg::with_name("state-cache-size")
                .long("state-cache-size")
                .value_name("SIZE")
                .help("Specifies how many states the database should cache in memory [default: 5]")
                .takes_value(true)
        )
        /*
         * Purge.
         */
        .arg(
            Arg::with_name("purge-db")
                .long("purge-db")
                .help("If present, the chain database will be deleted. Use with caution.")
        )
}
<commit_msg>Stop eth1 from always enabling (#1125)<commit_after>use clap::{App, Arg};

pub fn cli_app<'a, 'b>() -> App<'a, 'b> {
    App::new("beacon_node")
        .visible_aliases(&["b", "bn", "beacon"])
        .version(crate_version!())
        .author("Sigma Prime <contact@sigmaprime.io>")
        .about("The primary component which connects to the Ethereum 2.0 P2P network and \
                downloads, verifies and stores blocks. Provides a HTTP API for querying \
                the beacon chain and publishing messages to the network.")
        /*
         * Configuration directory locations.
         */
        .arg(
            Arg::with_name("network-dir")
                .long("network-dir")
                .value_name("DIR")
                .help("Data directory for network keys. Defaults to network/ inside the beacon node \
                       dir.")
                .takes_value(true)
        )
        .arg(
            Arg::with_name("freezer-dir")
                .long("freezer-dir")
                .value_name("DIR")
                .help("Data directory for the freezer database.")
                .takes_value(true)
        )
        /*
         * Network parameters.
         */
        .arg(
            Arg::with_name("zero-ports")
                .long("zero-ports")
                .short("z")
                .help("Sets all listening TCP/UDP ports to 0, allowing the OS to choose some \
                       arbitrary free ports.")
                .takes_value(false),
        )
        .arg(
            Arg::with_name("listen-address")
                .long("listen-address")
                .value_name("ADDRESS")
                .help("The address lighthouse will listen for UDP and TCP connections.")
                .default_value("0.0.0.0")
                .takes_value(true)
        )
        .arg(
            Arg::with_name("port")
                .long("port")
                .value_name("PORT")
                .help("The TCP/UDP port to listen on. The UDP port can be modified by the --discovery-port flag.")
                .default_value("9000")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("discovery-port")
                .long("discovery-port")
                .value_name("PORT")
                .help("The UDP port that discovery will listen on. Defaults to `port`")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("maxpeers")
                .long("maxpeers")
                .help("The maximum number of peers.")
                .default_value("50")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("boot-nodes")
                .long("boot-nodes")
                .allow_hyphen_values(true)
                .value_name("ENR-LIST")
                .help("One or more comma-delimited base64-encoded ENR's to bootstrap the p2p network.")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("enr-udp-port")
                .long("enr-udp-port")
                .value_name("PORT")
                .help("The UDP port of the local ENR. Set this only if you are sure other nodes can connect to your local node on this port.")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("enr-tcp-port")
                .long("enr-tcp-port")
                .value_name("PORT")
                .help("The TCP port of the local ENR. Set this only if you are sure other nodes can connect to your local node on this port.\
                    The --port flag is used if this is not set.")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("enr-address")
                .long("enr-address")
                .value_name("ADDRESS")
                .help("The IP address/ DNS address to broadcast to other peers on how to reach this node. \
                If a DNS address is provided, the enr-address is set to the IP address it resolves to and \
                does not auto-update based on PONG responses in discovery. \
                Set this only if you are sure other nodes can connect to your local node on this address. \
                Discovery will automatically find your external address,if possible.")
                .requires("enr-udp-port")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("enr-match")
                .short("e")
                .long("enr-match")
                .help("Sets the local ENR IP address and port to match those set for lighthouse. \
                Specifically, the IP address will be the value of --listen-address and the UDP port will be --discovery-port.")
        )
        .arg(
            Arg::with_name("disable-enr-auto-update")
                .short("x")
                .long("disable-enr-auto-update")
                .help("Discovery automatically updates the nodes local ENR with an external IP address and port as seen by other peers on the network. \
                This disables this feature, fixing the ENR's IP/PORT to those specified on boot.")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("libp2p-addresses")
                .long("libp2p-addresses")
                .value_name("MULTIADDR")
                .help("One or more comma-delimited multiaddrs to manually connect to a libp2p peer \
                       without an ENR.")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("p2p-priv-key")
                .long("p2p-priv-key")
                .value_name("HEX")
                .help("A secp256k1 secret key, represented as ASCII-encoded hex bytes (with or \
                       without 0x prefix). Default is either loaded from disk or generated \
                       automatically.")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("random-propagation")
                .long("random-propagation")
                .value_name("INTEGER")
                .takes_value(true)
                .help("Specifies (as a percentage) the likelihood of propagating blocks and \
                       attestations. This should only be used for testing networking elements. The \
                       value must like in the range 1-100. Default is 100.")
        )
        /* REST API related arguments */
        .arg(
            Arg::with_name("http")
                .long("http")
                .help("Enable RESTful HTTP API server. Disabled by default.")
                .takes_value(false),
        )
        .arg(
            Arg::with_name("http-address")
                .long("http-address")
                .value_name("ADDRESS")
                .help("Set the listen address for the RESTful HTTP API server.")
                .default_value("127.0.0.1")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("http-port")
                .long("http-port")
                .value_name("PORT")
                .help("Set the listen TCP port for the RESTful HTTP API server.")
                .default_value("5052")
                .takes_value(true),
        )
        /* Websocket related arguments */
        .arg(
            Arg::with_name("ws")
                .long("ws")
                .help("Enable the websocket server. Disabled by default.")
                .takes_value(false),
        )
        .arg(
            Arg::with_name("ws-address")
                .long("ws-address")
                .value_name("ADDRESS")
                .help("Set the listen address for the websocket server.")
                .default_value("127.0.0.1")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("ws-port")
                .long("ws-port")
                .value_name("PORT")
                .help("Set the listen TCP port for the websocket server.")
                .default_value("5053")
                .takes_value(true),
        )

        /*
         * Eth1 Integration
         */
        .arg(
            Arg::with_name("eth1")
                .long("eth1")
                .help("If present the node will connect to an eth1 node. This is required for \
                       block production, you must use this flag if you wish to serve a validator.")
                .takes_value(false),
        )
        .arg(
            Arg::with_name("dummy-eth1")
                .long("dummy-eth1")
                .conflicts_with("eth1")
                .help("If present, uses an eth1 backend that generates static dummy data.\
                      Identical to the method used at the 2019 Canada interop.")
        )
        .arg(
            Arg::with_name("eth1-endpoint")
                .long("eth1-endpoint")
                .value_name("HTTP-ENDPOINT")
                .help("Specifies the server for a web3 connection to the Eth1 chain. Also enables the --eth1 flag. Defaults to http://127.0.0.1:8545.")
                .takes_value(true)
        )
        .arg(
            Arg::with_name("slots-per-restore-point")
                .long("slots-per-restore-point")
                .value_name("SLOT_COUNT")
                .help("Specifies how often a freezer DB restore point should be stored. \
                       DO NOT DECREASE AFTER INITIALIZATION. [default: 2048 (mainnet) or 64 (minimal)]")
                .takes_value(true)
        )
        .arg(
            Arg::with_name("block-cache-size")
                .long("block-cache-size")
                .value_name("SIZE")
                .help("Specifies how many blocks the database should cache in memory [default: 5]")
                .takes_value(true)
        )
        .arg(
            Arg::with_name("state-cache-size")
                .long("state-cache-size")
                .value_name("SIZE")
                .help("Specifies how many states the database should cache in memory [default: 5]")
                .takes_value(true)
        )
        /*
         * Purge.
         */
        .arg(
            Arg::with_name("purge-db")
                .long("purge-db")
                .help("If present, the chain database will be deleted. Use with caution.")
        )
}
<|endoftext|>


<commit_before>//
// Copyright 2021 The Project Oak Authors
//
// Licensed 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.
//

// Required for enabling benchmark tests.
#![feature(test)]

use anyhow::Context;
use log::Level;
use prost::Message;
use serde_derive::Deserialize;
use std::{
    collections::HashMap,
    fs,
    net::{Ipv6Addr, SocketAddr},
    sync::{Arc, RwLock},
    time::{Duration, Instant},
};
use structopt::StructOpt;

mod logger;
mod server;
use crate::{logger::Logger, server::create_and_start_server};

#[cfg(test)]
mod tests;
#[derive(Deserialize, Debug)]
#[serde(deny_unknown_fields)]
struct Config {
    /// URL of a file to GET over HTTP containing key / value entries in protobuf binary format for
    /// lookup. If empty or not provided, no data is available for lookup.
    #[serde(default)]
    lookup_data_url: String,
    /// How often to refresh the lookup data. If not provided, data is only loaded once at startup.
    #[serde(with = "humantime_serde")]
    lookup_data_download_period: Option<Duration>,
}

/// Command line options for the Oak loader.
///
/// In general, when adding new configuration parameters, they should go in the `Config` struct
/// instead of here, and provided as part of the config TOML file by the developer, who would
/// normally bundle it with the Docker image of the Oak Functions Loader.
#[derive(StructOpt, Clone, Debug)]
#[structopt(about = "Oak Functions Loader")]
pub struct Opt {
    #[structopt(
        long,
        default_value = "8080",
        help = "Port number that the server listens on."
    )]
    http_listen_port: u16,
    #[structopt(
        long,
        help = "Path to a Wasm file to be loaded and executed per invocation. The Wasm module must export a function named `main`."
    )]
    wasm_path: String,
    #[structopt(
        long,
        help = "Path to a file containing configuration parameters in TOML format."
    )]
    config_path: String,
}

struct LookupData {
    lookup_data_url: String,
    entries: RwLock<HashMap<Vec<u8>, Vec<u8>>>,
    logger: Logger,
}

impl LookupData {
    /// Creates a new [`LookupData`] instance that can refresh its internal entries from the
    /// provided data file URL.
    ///
    /// Entries in the data file path must be consecutive binary encoded and length delimited
    /// protobuf messages according to the definition in `/oak_functions/proto/lookup_data.proto`.
    ///
    /// The returned instance is empty, and must be populated by calling the [`LookupData::refresh`]
    /// method at least once.
    fn new_empty(lookup_data_url: &str, logger: Logger) -> LookupData {
        LookupData {
            lookup_data_url: lookup_data_url.to_string(),
            entries: RwLock::new(HashMap::new()),
            logger,
        }
    }

    /// Refreshes the internal entries of this struct from the data file URL provided at
    /// construction time.
    ///
    /// If successful, entries are completely replaced (i.e. not merged).
    ///
    /// If there is any error while reading or parsing the data, an error is returned by this
    /// method, and existing entries are left untouched. The caller may retry the refresh operation
    /// at a future time.
    async fn refresh(&self) -> anyhow::Result<()> {
        // TODO(#1930): Avoid loading the entire file in memory for parsing.
        let client = hyper::Client::new();
        let res = client
            .get(
                self.lookup_data_url
                    .parse()
                    .context("could not parse lookup data URL")?,
            )
            .await
            .context("could not fetch lookup data")?;

        let start = Instant::now();
        let lookup_data_buf = hyper::body::to_bytes(res.into_body())
            .await
            .context("could not read lookup data response body")?;
        self.logger.log_public(
            Level::Debug,
            &format!("lookup data download time: {:?}", start.elapsed()),
        );

        let start = Instant::now();
        let entries = parse_lookup_entries(&mut lookup_data_buf.as_ref())
            .context("could not parse lookup data")?;

        self.logger.log_public(
            Level::Info,
            &format!("loaded {} entries of lookup data", entries.len()),
        );
        self.logger.log_public(
            Level::Debug,
            &format!("lookup data parsing time: {:?}", start.elapsed()),
        );

        // This block is here to emphasize and ensure that the write lock is only held for a very
        // short time.
        let start = Instant::now();
        {
            *self
                .entries
                .write()
                .expect("could not lock entries for write") = entries;
        }
        self.logger.log_public(
            Level::Debug,
            &format!(
                "lookup data write lock acquisition time: {:?}",
                start.elapsed()
            ),
        );

        Ok(())
    }

    /// Convenience getter for an individual entry that reduces lock contention by cloning the
    /// resulting value as quickly as possible and returning it instead of a reference.
    #[allow(dead_code)]
    fn get(&self, key: &[u8]) -> Option<Vec<u8>> {
        self.entries
            .read()
            .expect("could not lock entries for read")
            .get(key)
            .cloned()
    }

    #[allow(dead_code)]
    fn len(&self) -> usize {
        self.entries
            .read()
            .expect("could not lock entries for read")
            .len()
    }
}

fn parse_lookup_entries<B: bytes::Buf>(
    lookup_data_buffer: B,
) -> anyhow::Result<HashMap<Vec<u8>, Vec<u8>>> {
    let mut lookup_data_buffer = lookup_data_buffer;
    let mut entries = HashMap::new();
    while lookup_data_buffer.has_remaining() {
        let entry =
            oak_functions_abi::proto::Entry::decode_length_delimited(&mut lookup_data_buffer)
                .context("could not decode entry")?;
        entries.insert(entry.key, entry.value);
    }
    Ok(entries)
}

async fn background_refresh_lookup_data(lookup_data: &LookupData, period: Duration) {
    // Create an interval that starts after `period`, since the data was already refreshed
    // initially.
    let mut interval = tokio::time::interval_at(tokio::time::Instant::now() + period, period);
    loop {
        interval.tick().await;
        // If there is an error, we skip the current refresh and wait for the next tick.
        if let Err(err) = lookup_data.refresh().await {
            lookup_data.logger.log_public(
                Level::Error,
                &format!("error refreshing lookup data: {}", err),
            );
        }
    }
}

/// Main execution point for the Oak Functions Loader.
#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let opt = Opt::from_args();

    let config_file_bytes = fs::read(&opt.config_path)
        .with_context(|| format!("Couldn't read config file {}", &opt.config_path))?;
    let config: Config =
        toml::from_slice(&config_file_bytes).context("Couldn't parse config file")?;
    // Print the parsed config to STDERR regardless of whether the logger is initialized.

    // TODO(#1971): Make maximum log level configurable.
    let logger = Logger::default();

    logger.log_public(Level::Info, &format!("parsed config file:\n{:#?}", config));

    // For now the server runs in the same thread, so `notify_sender` is not really needed.
    let (_notify_sender, notify_receiver) = tokio::sync::oneshot::channel::<()>();

    let wasm_module_bytes = fs::read(&opt.wasm_path)
        .with_context(|| format!("Couldn't read Wasm file {}", &opt.wasm_path))?;

    let lookup_data = Arc::new(LookupData::new_empty(
        &config.lookup_data_url,
        logger.clone(),
    ));
    if !config.lookup_data_url.is_empty() {
        // First load the lookup data upfront in a blocking fashion.
        // TODO(#1930): Retry the initial lookup a few times if it fails.
        lookup_data
            .refresh()
            .await
            .context("Couldn't perform initial load of lookup data")?;
        if let Some(lookup_data_download_period) = config.lookup_data_download_period {
            // Create background task to periodically refresh the lookup data.
            let lookup_data = lookup_data.clone();
            tokio::spawn(async move {
                background_refresh_lookup_data(&lookup_data, lookup_data_download_period).await
            });
        };
    }

    // TODO(#1930): Pass lookup data to the server instance.

    // Start HTTP server.
    let address = SocketAddr::from((Ipv6Addr::UNSPECIFIED, opt.http_listen_port));
    create_and_start_server(&address, &wasm_module_bytes, notify_receiver, logger).await
}
<commit_msg>Unify logs for lookup data (#1988)<commit_after>//
// Copyright 2021 The Project Oak Authors
//
// Licensed 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.
//

// Required for enabling benchmark tests.
#![feature(test)]

use anyhow::Context;
use log::Level;
use prost::Message;
use serde_derive::Deserialize;
use std::{
    collections::HashMap,
    fs,
    net::{Ipv6Addr, SocketAddr},
    sync::{Arc, RwLock},
    time::{Duration, Instant},
};
use structopt::StructOpt;

mod logger;
mod server;
use crate::{logger::Logger, server::create_and_start_server};

#[cfg(test)]
mod tests;
#[derive(Deserialize, Debug)]
#[serde(deny_unknown_fields)]
struct Config {
    /// URL of a file to GET over HTTP containing key / value entries in protobuf binary format for
    /// lookup. If empty or not provided, no data is available for lookup.
    #[serde(default)]
    lookup_data_url: String,
    /// How often to refresh the lookup data. If not provided, data is only loaded once at startup.
    #[serde(with = "humantime_serde")]
    lookup_data_download_period: Option<Duration>,
}

/// Command line options for the Oak loader.
///
/// In general, when adding new configuration parameters, they should go in the `Config` struct
/// instead of here, and provided as part of the config TOML file by the developer, who would
/// normally bundle it with the Docker image of the Oak Functions Loader.
#[derive(StructOpt, Clone, Debug)]
#[structopt(about = "Oak Functions Loader")]
pub struct Opt {
    #[structopt(
        long,
        default_value = "8080",
        help = "Port number that the server listens on."
    )]
    http_listen_port: u16,
    #[structopt(
        long,
        help = "Path to a Wasm file to be loaded and executed per invocation. The Wasm module must export a function named `main`."
    )]
    wasm_path: String,
    #[structopt(
        long,
        help = "Path to a file containing configuration parameters in TOML format."
    )]
    config_path: String,
}

struct LookupData {
    lookup_data_url: String,
    entries: RwLock<HashMap<Vec<u8>, Vec<u8>>>,
    logger: Logger,
}

impl LookupData {
    /// Creates a new [`LookupData`] instance that can refresh its internal entries from the
    /// provided data file URL.
    ///
    /// Entries in the data file path must be consecutive binary encoded and length delimited
    /// protobuf messages according to the definition in `/oak_functions/proto/lookup_data.proto`.
    ///
    /// The returned instance is empty, and must be populated by calling the [`LookupData::refresh`]
    /// method at least once.
    fn new_empty(lookup_data_url: &str, logger: Logger) -> LookupData {
        LookupData {
            lookup_data_url: lookup_data_url.to_string(),
            entries: RwLock::new(HashMap::new()),
            logger,
        }
    }

    /// Refreshes the internal entries of this struct from the data file URL provided at
    /// construction time.
    ///
    /// If successful, entries are completely replaced (i.e. not merged).
    ///
    /// If there is any error while reading or parsing the data, an error is returned by this
    /// method, and existing entries are left untouched. The caller may retry the refresh operation
    /// at a future time.
    async fn refresh(&self) -> anyhow::Result<()> {
        // TODO(#1930): Avoid loading the entire file in memory for parsing.
        let client = hyper::Client::new();
        let res = client
            .get(
                self.lookup_data_url
                    .parse()
                    .context("could not parse lookup data URL")?,
            )
            .await
            .context("could not fetch lookup data")?;

        let start = Instant::now();
        let lookup_data_buf = hyper::body::to_bytes(res.into_body())
            .await
            .context("could not read lookup data response body")?;
        self.logger.log_public(
            Level::Info,
            &format!(
                "downloaded {} bytes of lookup data in {:.0?}",
                lookup_data_buf.len(),
                start.elapsed()
            ),
        );

        let start = Instant::now();
        let entries = parse_lookup_entries(&mut lookup_data_buf.as_ref())
            .context("could not parse lookup data")?;

        self.logger.log_public(
            Level::Info,
            &format!(
                "parsed {} entries of lookup data in {:.0?}",
                entries.len(),
                start.elapsed()
            ),
        );

        // This block is here to emphasize and ensure that the write lock is only held for a very
        // short time.
        let start = Instant::now();
        {
            *self
                .entries
                .write()
                .expect("could not lock entries for write") = entries;
        }
        self.logger.log_public(
            Level::Debug,
            &format!(
                "lookup data write lock acquisition time: {:.0?}",
                start.elapsed()
            ),
        );

        Ok(())
    }

    /// Convenience getter for an individual entry that reduces lock contention by cloning the
    /// resulting value as quickly as possible and returning it instead of a reference.
    #[allow(dead_code)]
    fn get(&self, key: &[u8]) -> Option<Vec<u8>> {
        self.entries
            .read()
            .expect("could not lock entries for read")
            .get(key)
            .cloned()
    }

    #[allow(dead_code)]
    fn len(&self) -> usize {
        self.entries
            .read()
            .expect("could not lock entries for read")
            .len()
    }
}

fn parse_lookup_entries<B: bytes::Buf>(
    lookup_data_buffer: B,
) -> anyhow::Result<HashMap<Vec<u8>, Vec<u8>>> {
    let mut lookup_data_buffer = lookup_data_buffer;
    let mut entries = HashMap::new();
    while lookup_data_buffer.has_remaining() {
        let entry =
            oak_functions_abi::proto::Entry::decode_length_delimited(&mut lookup_data_buffer)
                .context("could not decode entry")?;
        entries.insert(entry.key, entry.value);
    }
    Ok(entries)
}

async fn background_refresh_lookup_data(lookup_data: &LookupData, period: Duration) {
    // Create an interval that starts after `period`, since the data was already refreshed
    // initially.
    let mut interval = tokio::time::interval_at(tokio::time::Instant::now() + period, period);
    loop {
        interval.tick().await;
        // If there is an error, we skip the current refresh and wait for the next tick.
        if let Err(err) = lookup_data.refresh().await {
            lookup_data.logger.log_public(
                Level::Error,
                &format!("error refreshing lookup data: {}", err),
            );
        }
    }
}

/// Main execution point for the Oak Functions Loader.
#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let opt = Opt::from_args();

    let config_file_bytes = fs::read(&opt.config_path)
        .with_context(|| format!("Couldn't read config file {}", &opt.config_path))?;
    let config: Config =
        toml::from_slice(&config_file_bytes).context("Couldn't parse config file")?;
    // Print the parsed config to STDERR regardless of whether the logger is initialized.

    // TODO(#1971): Make maximum log level configurable.
    let logger = Logger::default();

    logger.log_public(Level::Info, &format!("parsed config file:\n{:#?}", config));

    // For now the server runs in the same thread, so `notify_sender` is not really needed.
    let (_notify_sender, notify_receiver) = tokio::sync::oneshot::channel::<()>();

    let wasm_module_bytes = fs::read(&opt.wasm_path)
        .with_context(|| format!("Couldn't read Wasm file {}", &opt.wasm_path))?;

    let lookup_data = Arc::new(LookupData::new_empty(
        &config.lookup_data_url,
        logger.clone(),
    ));
    if !config.lookup_data_url.is_empty() {
        // First load the lookup data upfront in a blocking fashion.
        // TODO(#1930): Retry the initial lookup a few times if it fails.
        lookup_data
            .refresh()
            .await
            .context("Couldn't perform initial load of lookup data")?;
        if let Some(lookup_data_download_period) = config.lookup_data_download_period {
            // Create background task to periodically refresh the lookup data.
            let lookup_data = lookup_data.clone();
            tokio::spawn(async move {
                background_refresh_lookup_data(&lookup_data, lookup_data_download_period).await
            });
        };
    }

    // TODO(#1930): Pass lookup data to the server instance.

    // Start HTTP server.
    let address = SocketAddr::from((Ipv6Addr::UNSPECIFIED, opt.http_listen_port));
    create_and_start_server(&address, &wasm_module_bytes, notify_receiver, logger).await
}
<|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/. */

//! Tracking of pending loads in a document.
//! https://html.spec.whatwg.org/multipage/#the-end

use dom::bindings::js::JS;
use dom::document::Document;
use ipc_channel::ipc::IpcSender;
use net_traits::{CoreResourceMsg, FetchResponseMsg, ResourceThreads, IpcSend};
use net_traits::request::RequestInit;
use servo_url::ServoUrl;
use std::thread;

#[derive(JSTraceable, PartialEq, Clone, Debug, HeapSizeOf)]
pub enum LoadType {
    Image(ServoUrl),
    Script(ServoUrl),
    Subframe(ServoUrl),
    Stylesheet(ServoUrl),
    PageSource(ServoUrl),
    Media(ServoUrl),
}

impl LoadType {
    fn url(&self) -> &ServoUrl {
        match *self {
            LoadType::Image(ref url) |
            LoadType::Script(ref url) |
            LoadType::Subframe(ref url) |
            LoadType::Stylesheet(ref url) |
            LoadType::Media(ref url) |
            LoadType::PageSource(ref url) => url,
        }
    }
}

/// Canary value ensuring that manually added blocking loads (ie. ones that weren't
/// created via DocumentLoader::fetch_async) are always removed by the time
/// that the owner is destroyed.
#[derive(JSTraceable, HeapSizeOf)]
#[must_root]
pub struct LoadBlocker {
    /// The document whose load event is blocked by this object existing.
    doc: JS<Document>,
    /// The load that is blocking the document's load event.
    load: Option<LoadType>,
}

impl LoadBlocker {
    /// Mark the document's load event as blocked on this new load.
    pub fn new(doc: &Document, load: LoadType) -> LoadBlocker {
        doc.mut_loader().add_blocking_load(load.clone());
        LoadBlocker {
            doc: JS::from_ref(doc),
            load: Some(load),
        }
    }

    /// Remove this load from the associated document's list of blocking loads.
    pub fn terminate(blocker: &mut Option<LoadBlocker>) {
        if let Some(this) = blocker.as_mut() {
            this.doc.finish_load(this.load.take().unwrap());
        }
        *blocker = None;
    }

    /// Return the url associated with this load.
    pub fn url(&self) -> Option<&ServoUrl> {
        self.load.as_ref().map(LoadType::url)
    }
}

impl Drop for LoadBlocker {
    fn drop(&mut self) {
        if !thread::panicking() {
            debug_assert!(self.load.is_none());
        }
    }
}

#[derive(JSTraceable, HeapSizeOf)]
pub struct DocumentLoader {
    resource_threads: ResourceThreads,
    blocking_loads: Vec<LoadType>,
    events_inhibited: bool,
}

impl DocumentLoader {
    pub fn new(existing: &DocumentLoader) -> DocumentLoader {
        DocumentLoader::new_with_threads(existing.resource_threads.clone(), None)
    }

    pub fn new_with_threads(resource_threads: ResourceThreads,
                            initial_load: Option<ServoUrl>) -> DocumentLoader {
        debug!("Initial blocking load {:?}.", initial_load);
        let initial_loads = initial_load.into_iter().map(LoadType::PageSource).collect();

        DocumentLoader {
            resource_threads: resource_threads,
            blocking_loads: initial_loads,
            events_inhibited: false,
        }
    }

    /// Add a load to the list of blocking loads.
    fn add_blocking_load(&mut self, load: LoadType) {
        debug!("Adding blocking load {:?} ({}).", load, self.blocking_loads.len());
        self.blocking_loads.push(load);
    }

    /// Initiate a new fetch.
    pub fn fetch_async(&mut self,
                       load: LoadType,
                       request: RequestInit,
                       fetch_target: IpcSender<FetchResponseMsg>) {
        self.add_blocking_load(load);
        self.fetch_async_background(request, fetch_target);
    }

    /// Initiate a new fetch that does not block the document load event.
    pub fn fetch_async_background(&mut self,
                                  request: RequestInit,
                                  fetch_target: IpcSender<FetchResponseMsg>) {
        self.resource_threads.sender().send(CoreResourceMsg::Fetch(request, fetch_target)).unwrap();
    }

    /// Mark an in-progress network request complete.
    pub fn finish_load(&mut self, load: &LoadType) {
        debug!("Removing blocking load {:?} ({}).", load, self.blocking_loads.len());
        let idx = self.blocking_loads.iter().position(|unfinished| *unfinished == *load);
        self.blocking_loads.remove(idx.expect(&format!("unknown completed load {:?}", load)));
    }

    pub fn is_blocked(&self) -> bool {
        // TODO: Ensure that we report blocked if parsing is still ongoing.
        !self.blocking_loads.is_empty()
    }

    pub fn inhibit_events(&mut self) {
        self.events_inhibited = true;
    }

    pub fn events_inhibited(&self) -> bool {
        self.events_inhibited
    }

    pub fn resource_threads(&self) -> &ResourceThreads {
        &self.resource_threads
    }
}
<commit_msg>Remove mutability requirement from fetch_async_background.<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/. */

//! Tracking of pending loads in a document.
//! https://html.spec.whatwg.org/multipage/#the-end

use dom::bindings::js::JS;
use dom::document::Document;
use ipc_channel::ipc::IpcSender;
use net_traits::{CoreResourceMsg, FetchResponseMsg, ResourceThreads, IpcSend};
use net_traits::request::RequestInit;
use servo_url::ServoUrl;
use std::thread;

#[derive(JSTraceable, PartialEq, Clone, Debug, HeapSizeOf)]
pub enum LoadType {
    Image(ServoUrl),
    Script(ServoUrl),
    Subframe(ServoUrl),
    Stylesheet(ServoUrl),
    PageSource(ServoUrl),
    Media(ServoUrl),
}

impl LoadType {
    fn url(&self) -> &ServoUrl {
        match *self {
            LoadType::Image(ref url) |
            LoadType::Script(ref url) |
            LoadType::Subframe(ref url) |
            LoadType::Stylesheet(ref url) |
            LoadType::Media(ref url) |
            LoadType::PageSource(ref url) => url,
        }
    }
}

/// Canary value ensuring that manually added blocking loads (ie. ones that weren't
/// created via DocumentLoader::fetch_async) are always removed by the time
/// that the owner is destroyed.
#[derive(JSTraceable, HeapSizeOf)]
#[must_root]
pub struct LoadBlocker {
    /// The document whose load event is blocked by this object existing.
    doc: JS<Document>,
    /// The load that is blocking the document's load event.
    load: Option<LoadType>,
}

impl LoadBlocker {
    /// Mark the document's load event as blocked on this new load.
    pub fn new(doc: &Document, load: LoadType) -> LoadBlocker {
        doc.mut_loader().add_blocking_load(load.clone());
        LoadBlocker {
            doc: JS::from_ref(doc),
            load: Some(load),
        }
    }

    /// Remove this load from the associated document's list of blocking loads.
    pub fn terminate(blocker: &mut Option<LoadBlocker>) {
        if let Some(this) = blocker.as_mut() {
            this.doc.finish_load(this.load.take().unwrap());
        }
        *blocker = None;
    }

    /// Return the url associated with this load.
    pub fn url(&self) -> Option<&ServoUrl> {
        self.load.as_ref().map(LoadType::url)
    }
}

impl Drop for LoadBlocker {
    fn drop(&mut self) {
        if !thread::panicking() {
            debug_assert!(self.load.is_none());
        }
    }
}

#[derive(JSTraceable, HeapSizeOf)]
pub struct DocumentLoader {
    resource_threads: ResourceThreads,
    blocking_loads: Vec<LoadType>,
    events_inhibited: bool,
}

impl DocumentLoader {
    pub fn new(existing: &DocumentLoader) -> DocumentLoader {
        DocumentLoader::new_with_threads(existing.resource_threads.clone(), None)
    }

    pub fn new_with_threads(resource_threads: ResourceThreads,
                            initial_load: Option<ServoUrl>) -> DocumentLoader {
        debug!("Initial blocking load {:?}.", initial_load);
        let initial_loads = initial_load.into_iter().map(LoadType::PageSource).collect();

        DocumentLoader {
            resource_threads: resource_threads,
            blocking_loads: initial_loads,
            events_inhibited: false,
        }
    }

    /// Add a load to the list of blocking loads.
    fn add_blocking_load(&mut self, load: LoadType) {
        debug!("Adding blocking load {:?} ({}).", load, self.blocking_loads.len());
        self.blocking_loads.push(load);
    }

    /// Initiate a new fetch.
    pub fn fetch_async(&mut self,
                       load: LoadType,
                       request: RequestInit,
                       fetch_target: IpcSender<FetchResponseMsg>) {
        self.add_blocking_load(load);
        self.fetch_async_background(request, fetch_target);
    }

    /// Initiate a new fetch that does not block the document load event.
    pub fn fetch_async_background(&self,
                                  request: RequestInit,
                                  fetch_target: IpcSender<FetchResponseMsg>) {
        self.resource_threads.sender().send(CoreResourceMsg::Fetch(request, fetch_target)).unwrap();
    }

    /// Mark an in-progress network request complete.
    pub fn finish_load(&mut self, load: &LoadType) {
        debug!("Removing blocking load {:?} ({}).", load, self.blocking_loads.len());
        let idx = self.blocking_loads.iter().position(|unfinished| *unfinished == *load);
        self.blocking_loads.remove(idx.expect(&format!("unknown completed load {:?}", load)));
    }

    pub fn is_blocked(&self) -> bool {
        // TODO: Ensure that we report blocked if parsing is still ongoing.
        !self.blocking_loads.is_empty()
    }

    pub fn inhibit_events(&mut self) {
        self.events_inhibited = true;
    }

    pub fn events_inhibited(&self) -> bool {
        self.events_inhibited
    }

    pub fn resource_threads(&self) -> &ResourceThreads {
        &self.resource_threads
    }
}
<|endoftext|>


<commit_before>use std::fmt;

#[derive(Debug, Clone)]
struct Tree {
    enter: Node,
    left: Option<Box<Tree>>,
    right: Option<Box<Tree>>,
}

/*impl fmt::Display for Tree {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "(this: {:#?}, \nleft: {:#?}, \nright: {:#?})",
            self.enter, self.left, self.right
        )
    }
}*/

//for red-black tree
#[derive(Debug, Clone, PartialEq)]
enum Color {
    Black,
    Red,
    White,
}

#[derive(Debug, Clone)]
struct Node {
    val: i32,
    color: Color,
}

impl Node {
    fn new(v: &i32) -> Node {
        Node {
            val: *v,
            color: Color::White,
        }
    }

    fn new_black(v: &i32) -> Node {
        let mut n = Node::new(v);
        n.make_black();
        n
    }
    fn new_red(v: &i32) -> Node {
        let mut n = Node::new(v);
        n.make_red();
        n
    }
    fn make_black(&mut self) {
        self.color = Color::Black;
    }

    fn make_red(&mut self) {
        self.color = Color::Red;
    }

    fn change_color(&mut self) {
        if self.color == Color::Black {
            self.color = Color::Red
        } else {
            self.color = Color::Black
        }
    }
}

type ReTr<'a> = Result<&'a mut Tree, String>;

impl Tree {
    fn new(v: &i32) -> Tree {
        let mut root = Node::new(v);
        root.make_black();
        Tree {
            enter: root,
            left: None,
            right: None,
        }
    }

    //fix color at ending of insert method
    fn insert(&mut self, v: &i32) {
        if *v == self.enter.val {
            println!("{}", "duplicate value");
            return;
        }

        {
            //need drop target which hold self ownship..
            //we need this ownship to fix color at endding of this method.
            let target = if *v > self.enter.val {
                &mut self.right
            } else {
                &mut self.left
            };

            match target {
                Some(next) => next.insert(v),
                None => {
                    let mut new_node = Node::new(v);
                    new_node.make_red();
                    let new_tree = Tree {
                        enter: new_node,
                        left: None,
                        right: None,
                    };
                    *target = Some(Box::new(new_tree));
                }
            }
        }

        //self.fix(v);
    }

    fn look_up_recursive(&mut self, v: &i32) -> ReTr {
        if *v == self.enter.val {
            return Ok(self);
        } else if *v > self.enter.val {
            if let Some(_) = self.right {
                let next = self.right.as_mut().unwrap();
                return next.look_up(v);
            } else {
                Err("not find".to_string())
            }
        } else {
            if let Some(_) = self.left {
                let next = self.left.as_mut().unwrap();
                return next.look_up(v);
            } else {
                Err("not find".to_string())
            }
        }
    }

    fn look_up(&mut self, v: &i32) -> ReTr {
        let mut this: &mut Self = self;
        loop {
            if *v == this.enter.val {
                return Ok(this);
            } else if *v > this.enter.val {
                if let None = this.right {
                    return Err("not find".to_string());
                }
                this = { this }.right.as_mut().unwrap();
            } else {
                if let None = this.left {
                    return Err("not find".to_string());
                }
                this = { this }.left.as_mut().unwrap();
            }
        }
    }

    fn look_up_father(&mut self, v: &i32) -> (ReTr, &'static str) {
        let mut this: &mut Self = self;
        let mut result: &str = "None";
        let mut next: &str = "None";
        loop {
            if *v == this.enter.val {
                return (Err("no father".to_string()), "this");
            }

            match (this.right.as_ref(), this.left.as_ref()) {
                (Some(r), Some(l)) => {
                    if r.enter.val == *v {
                        result = "right";
                    } else if l.enter.val == *v {
                        result = "left"
                    }
                }
                (Some(r), None) => {
                    if r.enter.val == *v {
                        result = "right";
                    }
                }
                (None, Some(l)) => {
                    if l.enter.val == *v {
                        result = "left";
                    }
                }
                _ => (),
            }

            if *v > this.enter.val {
                if let None = this.right {
                    return (Err("not find".to_string()), "Error");
                }
                next = "right";
            } else {
                if let None = this.left {
                    return (Err("not find".to_string()), "Error");
                }
                next = "left";
            }

            if result != "None" {
                return (Ok(this), result);
            } else if next == "right" {
                this = { this }.right.as_mut().unwrap();
            } else if next == "left" {
                this = { this }.left.as_mut().unwrap();
            }
        }
    }

    //find father, grandfather, and uncle
    fn look_up_family(
        &mut self,
        v: &i32,
    ) -> (
        Result<i32, &'static str>,
        Result<i32, &'static str>,
        Result<i32, &'static str>,
    ) {
        let this: &mut Self = self;
        let grandfather: i32;
        let father: i32;

        match this.look_up_father(v).0 {
            Ok(n) => {
                father = n.enter.val;
            }
            Err(_) => {
                return (
                    Err("father not find"),
                    Err("grandfather not find"),
                    Err("uncle not find"),
                )
            }
        }

        let uncle_option: &Option<Box<Tree>>;
        match this.look_up_father(&father) {
            (Ok(n), s) => {
                grandfather = n.enter.val;
                match s {
                    "left" => uncle_option = &n.right,
                    _ => uncle_option = &n.left,
                };
            }
            (Err(_), _) => {
                return (
                    Ok(father),
                    Err("grandfather not find"),
                    Err("uncle not find"),
                )
            }
        }

        match uncle_option {
            Some(un) => return (Ok(father), Ok(grandfather), Ok(un.enter.val)),
            _ => return (Ok(father), Ok(grandfather), Err("uncle not find")),
        }
    }

    fn cut_left_tree(&self) -> Option<Box<Self>> {
        self.left.clone()
    }

    fn cut_right_tree(&self) -> Option<Box<Self>> {
        self.right.clone()
    }

    fn left_shift(&mut self, v: &i32) {
        match self.look_up_father(v).0 {
            Ok(r) => {
                let father = r;
                let mut this = father.cut_left_tree().unwrap();
                let mut son = this.cut_right_tree().unwrap();
                let grand_son = son.cut_left_tree().unwrap();

                this.right = Some(grand_son);
                son.left = Some(this);
                father.left = Some(son);
                return;
            }
            _ => (),
        }

        let left = self.cut_left_tree().unwrap();
        let mut new_left = Tree::new(&self.enter.val);
        new_left.left = Some(left);

        let right = self.cut_right_tree().unwrap();
        let right_left = right.cut_left_tree().unwrap();
        let right_right = right.cut_right_tree().unwrap();
        new_left.right = Some(right_left);

        self.enter = Node::new_black(&right.enter.val);

        self.right = Some(right_right);

        self.left = Some(Box::new(new_left));
    }

    fn right_shift(&mut self, v: &i32) {
        match self.look_up_father(v).0 {
            Ok(r) => {
                let father = r;
                let mut this = father.cut_right_tree().unwrap();
                let mut son = this.cut_left_tree().unwrap();
                let grand_son = son.cut_right_tree();

                this.left = grand_son;
                son.right = Some(this);
                father.right = Some(son);
                return;
            }
            _ => (),
        }

        let right = self.cut_right_tree();
        let left = self.cut_left_tree().unwrap();
        let left_left = left.cut_left_tree().unwrap();
        let left_right = left.cut_right_tree().unwrap();

        let mut new_right = Tree::new(&self.enter.val);
        new_right.left = Some(left_right);
        new_right.right = right;

        self.enter = Node::new_black(&left.enter.val);
        self.left = Some(left_left);
        self.right = Some(Box::new(new_right));
    }

    //this fix method might not very efficiency because look_up_father..
    //always look from top of tree
    fn fix(&mut self, v: &i32) {
        //step 1
        let mut this: i32 = 0;
        match self.look_up_family(v) {
            (Ok(father), Ok(grandfather), Ok(uncle)) => {
                if let Ok(fff) = self.look_up_father(&grandfather).0 {
                    this = fff.enter.val;
                } else {
                    println!("{}", "no great-grandfather");
                    return;
                }

                if let Ok(ff) = self.look_up(&father) {
                    ff.enter.make_black();
                }

                if let Ok(un) = self.look_up(&uncle) {
                    un.enter.make_black();
                }

                if let Ok(ff) = self.look_up(&grandfather) {
                    ff.enter.make_red();
                }
            }
            (a, b, c) => println!("{:#?}, {:#?}, {:#?}", a, b, c),
        }

        //step 2
        self.left_shift(&this);
        this = self
            .look_up(&{ this })
            .unwrap()
            .left
            .as_mut()
            .unwrap()
            .enter
            .val;

        //stap3
        match self.look_up_family(&this) {
            (Ok(father), Ok(grandfather), _) => {
                if let Ok(ff) = self.look_up(&father) {
                    ff.enter.make_black();
                }

                if let Ok(gf) = self.look_up(&grandfather) {
                    gf.enter.make_red();
                    this = gf.enter.val;
                }
            }
            _ => (),
        }

        this = self.look_up_father(&this).0.unwrap().enter.val;
        self.right_shift(&this);
        self.enter.make_black();
    }
}

fn main() {
    //let mut a = Tree::new(&3);
    //a.insert(&2);
    //a.insert(&1);
    //a.insert(&4);
    //a.insert(&5);

    //println!("{:?}", a);
    //println!("{:?}", a.look_up(&1));
    //println!("{:?}", a.look_up_father(&4));
    //println!("{:?}", a);

    let mut b = Tree::new(&10);
    b.insert(&2);
    b.insert(&1);
    b.insert(&5);
    b.insert(&3);
    b.insert(&6);
    b.insert(&11);

    //println!("{:?}", b);
    //println!("{:?}", b.look_up_father(&2));
    let mut c = b.cut_left_tree().unwrap();
    //println!("{:#?}", c);
    c.left_shift(&2);
    println!("{:#?}", c);

    //let mut c = Tree::new(&2);
    //c.insert(&8);
    //c.insert(&1);
    //c.insert(&5);
    //c.insert(&4);
    //c.insert(&7);
    //c.insert(&9);
    //println!("{:?}", c.look_up_father(&8));
    //c.right_shift(&8);
    //println!("{:#?}", c);
}
<commit_msg>change insert method and a bit test cases in main<commit_after>use std::fmt;

#[derive(Debug, Clone)]
struct Tree {
    enter: Node,
    left: Option<Box<Tree>>,
    right: Option<Box<Tree>>,
}

/*impl fmt::Display for Tree {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "(this: {:#?}, \nleft: {:#?}, \nright: {:#?})",
            self.enter, self.left, self.right
        )
    }
}*/

//for red-black tree
#[derive(Debug, Clone, PartialEq)]
enum Color {
    Black,
    Red,
    White,
}

#[derive(Debug, Clone)]
struct Node {
    val: i32,
    color: Color,
}

impl Node {
    fn new(v: &i32) -> Node {
        Node {
            val: *v,
            color: Color::White,
        }
    }

    fn new_black(v: &i32) -> Node {
        let mut n = Node::new(v);
        n.make_black();
        n
    }
    fn new_red(v: &i32) -> Node {
        let mut n = Node::new(v);
        n.make_red();
        n
    }
    fn make_black(&mut self) {
        self.color = Color::Black;
    }

    fn make_red(&mut self) {
        self.color = Color::Red;
    }

    fn change_color(&mut self) {
        if self.color == Color::Black {
            self.color = Color::Red
        } else {
            self.color = Color::Black
        }
    }
}

type ReTr<'a> = Result<&'a mut Tree, String>;

impl Tree {
    fn new(v: &i32) -> Tree {
        let mut root = Node::new(v);
        root.make_black();
        Tree {
            enter: root,
            left: None,
            right: None,
        }
    }

    //fix color at ending of insert method
    fn insert_recursive(&mut self, v: &i32) {
        if *v == self.enter.val {
            println!("{}", "duplicate value");
            return;
        }

        let color_this: Color;
        //need drop target which hold self ownship..
        //we need this ownship to fix color at endding of this method.
        {
            let target = if *v > self.enter.val {
                &mut self.right
            } else {
                &mut self.left
            };

            match target {
                Some(next) => next.insert(v),
                None => {
                    let mut new_node = Node::new(v);
                    new_node.make_red();
                    let new_tree = Tree {
                        enter: new_node,
                        left: None,
                        right: None,
                    };
                    *target = Some(Box::new(new_tree));
                }
            }
        }
    }

    fn insert(&mut self, v: &i32) {
        let mut color: Color;

        {
            let mut this: &mut Self = self;
            loop {
                if *v == this.enter.val {
                    println!("{}", "duplicate value");
                    return;
                }

                if *v > this.enter.val {
                    if let Some(_) = this.right {
                        this = { this }.right.as_mut().unwrap();
                        continue;
                    } else {
                        let mut new_node = Node::new(v);
                        new_node.make_red();
                        let new_tree = Tree {
                            enter: new_node,
                            left: None,
                            right: None,
                        };
                        this.right = Some(Box::new(new_tree));
                        color = this.enter.color.clone();
                        break;
                    }
                }

                if *v < this.enter.val {
                    if let Some(_) = this.left {
                        this = { this }.left.as_mut().unwrap();
                        continue;
                    } else {
                        let mut new_node = Node::new(v);
                        new_node.make_red();
                        let new_tree = Tree {
                            enter: new_node,
                            left: None,
                            right: None,
                        };
                        this.left = Some(Box::new(new_tree));
                        color = this.enter.color.clone();
                        break;
                    }
                }
            }
        }

        if color == Color::Red {
            println!("{}", "fix here");
            //self.fix(v);
        }
    }

    fn look_up_recursive(&mut self, v: &i32) -> ReTr {
        if *v == self.enter.val {
            return Ok(self);
        } else if *v > self.enter.val {
            if let Some(_) = self.right {
                let next = self.right.as_mut().unwrap();
                return next.look_up(v);
            } else {
                Err("not find".to_string())
            }
        } else {
            if let Some(_) = self.left {
                let next = self.left.as_mut().unwrap();
                return next.look_up(v);
            } else {
                Err("not find".to_string())
            }
        }
    }

    fn look_up(&mut self, v: &i32) -> ReTr {
        let mut this: &mut Self = self;
        loop {
            if *v == this.enter.val {
                return Ok(this);
            } else if *v > this.enter.val {
                if let None = this.right {
                    return Err("not find".to_string());
                }
                this = { this }.right.as_mut().unwrap();
            } else {
                if let None = this.left {
                    return Err("not find".to_string());
                }
                this = { this }.left.as_mut().unwrap();
            }
        }
    }

    fn look_up_father(&mut self, v: &i32) -> (ReTr, &'static str) {
        let mut this: &mut Self = self;
        let mut result: &str = "None";
        let mut next: &str = "None";
        loop {
            if *v == this.enter.val {
                return (Err("no father".to_string()), "this");
            }

            match (this.right.as_ref(), this.left.as_ref()) {
                (Some(r), Some(l)) => {
                    if r.enter.val == *v {
                        result = "right";
                    } else if l.enter.val == *v {
                        result = "left"
                    }
                }
                (Some(r), None) => {
                    if r.enter.val == *v {
                        result = "right";
                    }
                }
                (None, Some(l)) => {
                    if l.enter.val == *v {
                        result = "left";
                    }
                }
                _ => (),
            }

            if *v > this.enter.val {
                if let None = this.right {
                    return (Err("not find".to_string()), "Error");
                }
                next = "right";
            } else {
                if let None = this.left {
                    return (Err("not find".to_string()), "Error");
                }
                next = "left";
            }

            if result != "None" {
                return (Ok(this), result);
            } else if next == "right" {
                this = { this }.right.as_mut().unwrap();
            } else if next == "left" {
                this = { this }.left.as_mut().unwrap();
            }
        }
    }

    //find father, grandfather, and uncle
    fn look_up_family(
        &mut self,
        v: &i32,
    ) -> (
        Result<i32, &'static str>,
        Result<i32, &'static str>,
        Result<i32, &'static str>,
    ) {
        let this: &mut Self = self;
        let grandfather: i32;
        let father: i32;

        match this.look_up_father(v).0 {
            Ok(n) => {
                father = n.enter.val;
            }
            Err(_) => {
                return (
                    Err("father not find"),
                    Err("grandfather not find"),
                    Err("uncle not find"),
                )
            }
        }

        let uncle_option: &Option<Box<Tree>>;
        match this.look_up_father(&father) {
            (Ok(n), s) => {
                grandfather = n.enter.val;
                match s {
                    "left" => uncle_option = &n.right,
                    _ => uncle_option = &n.left,
                };
            }
            (Err(_), _) => {
                return (
                    Ok(father),
                    Err("grandfather not find"),
                    Err("uncle not find"),
                )
            }
        }

        match uncle_option {
            Some(un) => return (Ok(father), Ok(grandfather), Ok(un.enter.val)),
            _ => return (Ok(father), Ok(grandfather), Err("uncle not find")),
        }
    }

    fn cut_left_tree(&self) -> Option<Box<Self>> {
        self.left.clone()
    }

    fn cut_right_tree(&self) -> Option<Box<Self>> {
        self.right.clone()
    }

    fn left_shift(&mut self, v: &i32) {
        match self.look_up_father(v).0 {
            Ok(r) => {
                let father = r;
                let mut this = father.cut_left_tree().unwrap();
                let mut son = this.cut_right_tree().unwrap();
                let grand_son = son.cut_left_tree().unwrap();

                this.right = Some(grand_son);
                son.left = Some(this);
                father.left = Some(son);
                return;
            }
            _ => (),
        }

        let left = self.cut_left_tree().unwrap();
        let mut new_left = Tree::new(&self.enter.val);
        new_left.left = Some(left);

        let right = self.cut_right_tree().unwrap();
        let right_left = right.cut_left_tree().unwrap();
        let right_right = right.cut_right_tree().unwrap();
        new_left.right = Some(right_left);

        self.enter = Node::new_black(&right.enter.val);

        self.right = Some(right_right);

        self.left = Some(Box::new(new_left));
    }

    fn right_shift(&mut self, v: &i32) {
        match self.look_up_father(v).0 {
            Ok(r) => {
                let father = r;
                let mut this = father.cut_right_tree().unwrap();
                let mut son = this.cut_left_tree().unwrap();
                let grand_son = son.cut_right_tree();

                this.left = grand_son;
                son.right = Some(this);
                father.right = Some(son);
                return;
            }
            _ => (),
        }

        let right = self.cut_right_tree();
        let left = self.cut_left_tree().unwrap();
        let left_left = left.cut_left_tree().unwrap();
        let left_right = left.cut_right_tree().unwrap();

        let mut new_right = Tree::new(&self.enter.val);
        new_right.left = Some(left_right);
        new_right.right = right;

        self.enter = Node::new_black(&left.enter.val);
        self.left = Some(left_left);
        self.right = Some(Box::new(new_right));
    }

    //this fix method might not very efficiency because look_up_father..
    //always look from top of tree
    fn fix(&mut self, v: &i32) {
        //step 1
        let mut this: i32 = 0;
        match self.look_up_family(v) {
            (Ok(father), Ok(grandfather), Ok(uncle)) => {
                if let Ok(fff) = self.look_up_father(&grandfather).0 {
                    this = fff.enter.val;
                } else {
                    println!("{}", "no great-grandfather");
                    return;
                }

                if let Ok(ff) = self.look_up(&father) {
                    ff.enter.make_black();
                }

                if let Ok(un) = self.look_up(&uncle) {
                    un.enter.make_black();
                }

                if let Ok(ff) = self.look_up(&grandfather) {
                    ff.enter.make_red();
                }
            }
            (a, b, c) => println!("{:#?}, {:#?}, {:#?}", a, b, c),
        }

        //step 2
        self.left_shift(&this);
        this = self
            .look_up(&{ this })
            .unwrap()
            .left
            .as_mut()
            .unwrap()
            .enter
            .val;

        //stap3
        match self.look_up_family(&this) {
            (Ok(father), Ok(grandfather), _) => {
                if let Ok(ff) = self.look_up(&father) {
                    ff.enter.make_black();
                }

                if let Ok(gf) = self.look_up(&grandfather) {
                    gf.enter.make_red();
                    this = gf.enter.val;
                }
            }
            _ => (),
        }

        this = self.look_up_father(&this).0.unwrap().enter.val;
        self.right_shift(&this);
        self.enter.make_black();
    }
}

fn main() {
    //let mut a = Tree::new(&3);
    //a.insert(&2);
    //a.insert(&1);
    //a.insert(&4);
    //a.insert(&5);

    //println!("{:?}", a);
    //println!("{:?}", a.look_up(&1));
    //println!("{:?}", a.look_up_father(&4));
    //println!("{:?}", a);

    //let mut b = Tree::new(&10);
    //b.insert(&2);
    //b.insert(&1);
    //b.insert(&5);
    //b.insert(&3);
    //b.insert(&6);
    //b.insert(&11);

    //println!("{:#?}", b);
    //println!("{:?}", b.look_up_father(&2));
    //let mut c = b.cut_left_tree().unwrap();
    //println!("{:#?}", c);
    //c.left_shift(&2);
    //println!("{:#?}", c);
    //println!("{:#?}", b.look_up_family(&1));
    //println!("{:#?}", b.look_up_family(&2));
    //println!("{:#?}", b.look_up_family(&10));

    let mut c = Tree::new(&11);
    c.insert(&2);
    c.insert(&14);
    c.insert(&1);
    c.insert(&7);
    c.insert(&13);
    c.insert(&8);
    c.insert(&5);
    c.insert(&4);
    //c.fix(&4);
    //println!("{:#?}", c.look_up_father(&4));
    //c.left_shift(&2);
    println!("{:#?}", c);
}
<|endoftext|>