<commit_before>use time::{at_utc, Tm, Timespec};
use json::{FromJsonnable, ToJsonnable};
use serde::de::{Error, Type};
use serde_json::value::Value;
use serde_json::error::Error as JsonError;
use serde_json::builder::ObjectBuilder;

impl FromJsonnable for Tm {
	// Deserialize via Timespec
	fn from_json(json: Value) -> Result<Self, JsonError> {
		match json {
			Value::Object(map) => {
				let sec =
					match map.get("sec") {
						Some(sec) =>
							match sec {
								&Value::I64(sec) => sec,
								&Value::U64(sec) => sec as i64,  // The types get weird here
								_ => return Err(JsonError::type_mismatch(Type::I64)),
							},
						None => return Err(JsonError::missing_field("Missing \"sec\"")),
					};
				let nsec =
					match map.get("nsec") {
						Some(nsec) =>
							match nsec {
								&Value::I64(nsec) => nsec as i32,
								&Value::U64(nsec) => nsec as i32,
								_ => return Err(JsonError::type_mismatch(Type::I32)),
							},
						None => return Err(JsonError::missing_field("Missing \"nsec\"")),
					};

				Ok(at_utc(Timespec::new(sec, nsec)))
			},
			_ => Err(JsonError::type_mismatch(Type::Struct)),
		}
	}
}

impl ToJsonnable for Tm {
	// Serialize via Timespec
	fn to_json(&self) -> Value {
		let spec = self.to_timespec();
		ObjectBuilder::new().insert("sec", &spec.sec)
		                    .insert("nsec", &spec.nsec)
		                    .unwrap()
	}
}
<commit_msg>Implement {From,To}Jsonnable for Vec<T><commit_after>use time::{at_utc, Tm, Timespec};
use json::{FromJsonnable, ToJsonnable};
use serde::de::{Error, Type};
use serde_json::value::Value;
use serde_json::error::Error as JsonError;
use serde_json::builder::ObjectBuilder;

impl FromJsonnable for Tm {
	// Deserialize via Timespec
	fn from_json(json: Value) -> Result<Self, JsonError> {
		match json {
			Value::Object(map) => {
				let sec =
					match map.get("sec") {
						Some(sec) =>
							match sec {
								&Value::I64(sec) => sec,
								&Value::U64(sec) => sec as i64,  // The types get weird here
								_ => return Err(JsonError::type_mismatch(Type::I64)),
							},
						None => return Err(JsonError::missing_field("Missing \"sec\"")),
					};
				let nsec =
					match map.get("nsec") {
						Some(nsec) =>
							match nsec {
								&Value::I64(nsec) => nsec as i32,
								&Value::U64(nsec) => nsec as i32,
								_ => return Err(JsonError::type_mismatch(Type::I32)),
							},
						None => return Err(JsonError::missing_field("Missing \"nsec\"")),
					};

				Ok(at_utc(Timespec::new(sec, nsec)))
			},
			_ => Err(JsonError::type_mismatch(Type::Struct)),
		}
	}
}

impl ToJsonnable for Tm {
	// Serialize via Timespec
	fn to_json(&self) -> Value {
		let spec = self.to_timespec();
		ObjectBuilder::new().insert("sec", &spec.sec)
		                    .insert("nsec", &spec.nsec)
		                    .unwrap()
	}
}

impl<T: FromJsonnable> FromJsonnable for Vec<T> {
	fn from_json(json: Value) -> Result<Self, JsonError> {
		match json {
			Value::Array(arr) => {
				let mut elems: Vec<T> = Vec::with_capacity(arr.len());
				for elem in arr {
					match T::from_json(elem) {
						Ok(elem) => elems.push(elem),
						Err(e) => return Err(e),
					}
				}
				Ok(elems)
			},
			_ => Err(JsonError::type_mismatch(Type::Seq)),
		}
	}
}

impl<T: ToJsonnable> ToJsonnable for Vec<T> {
	fn to_json(&self) -> Value {
		Value::Array(self.iter().map(|ref elem| elem.to_json()).collect())
	}
}
<|endoftext|>


<commit_before>//! A module for primitive rectangle drawing

pub use super::vga::Color;

use core::fmt::{Write,Error};

use super::vga::VGA;

/// Ann abstraction for drawing text inside a bounded box
///
/// A rectangle has boundaries. Drawing outside those boundaries will
/// not render any text on the screen.
///
/// Furthermore, a rectangle can draw a string, rather than just a single
/// character.
#[derive(Copy, Clone)]
pub struct Rectangle {
    height: usize,
    width: usize,
    pos: (usize, usize),
    cursor: (usize, usize), // position relative to the rectangle
    vga: VGA,
}

impl Rectangle {
    /// Create a new Rectangle.
    pub fn new(w: usize, h: usize, (row, col): (usize, usize)) -> Rectangle {
        Rectangle {
            height: h,
            width: w,
            pos: (row, col),
            cursor: (row, col),
            vga: VGA::new(),
        }
    }

    /// A static method to clear the screen.
    pub fn clear_screen() {
        let mut buff = VGA::new();
        buff.set_fg(Color::White);
        buff.set_bg(Color::Black);
        buff.clear_screen();
    }

    /// Paint the rectangle with the background color
    pub fn paint(&mut self) {
        self.vga.fill_rect(self.pos, self.height, self.width);
    }

    /// Draw the character at the cursor relative to the corner of rectangle
    /// if the character is inside the rectangle.
    pub fn put_char(&mut self, ch: char) {
        let (cr, cc) = self.cursor;
        if cr < self.height && cc < self.width {
            self.vga.put_char(ch);
        }
    }

    /// Render the string at the cursor using word wrapping.
    pub fn put_str(&mut self, s: &str) {
        for word in s.split(" ") {
            // word wrapping
            let len = word.len();
            let (mut cr, mut cc) = self.cursor;

            if cc + len >= self.width && len < self.width {
                cr += 1;
                cc = 0;
                self.set_cursor((cr, cc));
            }

            // draw a space before the word
            if cc > 0 {
                self.put_char(' ');

                // advance
                if cc + 1 >= self.width {
                    cr += 1;
                    cc = 0;
                } else {
                    cc += 1;
                }
                self.set_cursor((cr,cc));
            }

            for ch in word.chars() {
                // take char of special characters
                match ch {
                    '\n' => {
                        // don't draw a new line, just move the cursor
                        cr += 1;
                        cc = 0;
                    },
                    _ => {
                        // draw char
                        self.put_char(ch);

                        // advance
                        if cc + 1 >= self.width {
                            cr += 1;
                            cc = 0;
                        } else {
                            cc += 1;
                        }
                    }
                }

                // move cursor
                self.set_cursor((cr, cc));
            }
        }
    }

    /// Set the foreground color
    pub fn set_fg(&mut self, color: Color) {
        self.vga.set_fg(color);
    }

    /// Set the background color
    pub fn set_bg(&mut self, color: Color) {
        self.vga.set_bg(color);
    }

    /// Set the cursor position
    pub fn set_cursor(&mut self, (crow, ccol): (usize, usize)) {
        let (row,col) = self.pos;

        self.cursor = (crow, ccol);
        self.vga.set_cursor((row+crow, col+ccol));
    }
}

impl Write for Rectangle {
    fn write_str(&mut self, s: &str) -> Result<(), Error> {
        self.put_str(s);
        Result::Ok(())
    }
}
<commit_msg>Fix annoying bug that prints extra spaces<commit_after>//! A module for primitive rectangle drawing

pub use super::vga::Color;

use core::fmt::{Write,Error};

use super::vga::VGA;

/// Ann abstraction for drawing text inside a bounded box
///
/// A rectangle has boundaries. Drawing outside those boundaries will
/// not render any text on the screen.
///
/// Furthermore, a rectangle can draw a string, rather than just a single
/// character.
#[derive(Copy, Clone)]
pub struct Rectangle {
    height: usize,
    width: usize,
    pos: (usize, usize),
    cursor: (usize, usize), // position relative to the rectangle
    vga: VGA,
}

impl Rectangle {
    /// Create a new Rectangle.
    pub fn new(w: usize, h: usize, (row, col): (usize, usize)) -> Rectangle {
        Rectangle {
            height: h,
            width: w,
            pos: (row, col),
            cursor: (row, col),
            vga: VGA::new(),
        }
    }

    /// A static method to clear the screen.
    pub fn clear_screen() {
        let mut buff = VGA::new();
        buff.set_fg(Color::White);
        buff.set_bg(Color::Black);
        buff.clear_screen();
    }

    /// Paint the rectangle with the background color
    pub fn paint(&mut self) {
        self.vga.fill_rect(self.pos, self.height, self.width);
    }

    /// Draw the character at the cursor relative to the corner of rectangle
    /// if the character is inside the rectangle.
    pub fn put_char(&mut self, ch: char) {
        let (cr, cc) = self.cursor;
        if cr < self.height && cc < self.width {
            self.vga.put_char(ch);
        }
    }

    /// Render the string at the cursor using word wrapping.
    pub fn put_str(&mut self, s: &str) {
        let mut first = true;
        for word in s.split(" ") {
            // word wrapping
            let len = word.len();
            let (mut cr, mut cc) = self.cursor;

            if cc + len >= self.width && len < self.width {
                cr += 1;
                cc = 0;
                self.set_cursor((cr, cc));
            }

            // draw a space before the word
            if cc > 0 && !first {
                self.put_char(' ');

                // advance
                if cc + 1 >= self.width {
                    cr += 1;
                    cc = 0;
                } else {
                    cc += 1;
                }
                self.set_cursor((cr,cc));
            }

            for ch in word.chars() {
                // take char of special characters
                match ch {
                    '\n' => {
                        // don't draw a new line, just move the cursor
                        cr += 1;
                        cc = 0;
                    },
                    _ => {
                        // draw char
                        self.put_char(ch);

                        // advance
                        if cc + 1 >= self.width {
                            cr += 1;
                            cc = 0;
                        } else {
                            cc += 1;
                        }
                    }
                }

                // move cursor
                self.set_cursor((cr, cc));

                first = false;
            }
        }
    }

    /// Set the foreground color
    pub fn set_fg(&mut self, color: Color) {
        self.vga.set_fg(color);
    }

    /// Set the background color
    pub fn set_bg(&mut self, color: Color) {
        self.vga.set_bg(color);
    }

    /// Set the cursor position
    pub fn set_cursor(&mut self, (crow, ccol): (usize, usize)) {
        let (row,col) = self.pos;

        self.cursor = (crow, ccol);
        self.vga.set_cursor((row+crow, col+ccol));
    }
}

impl Write for Rectangle {
    fn write_str(&mut self, s: &str) -> Result<(), Error> {
        self.put_str(s);
        Result::Ok(())
    }
}
<|endoftext|>


<commit_before>#[doc (hidden)]
#[macro_export]
macro_rules! time_steward_define_simple_timeline {
  () => {
pub mod simple_timeline {
use std::collections::BTreeSet;
use std::mem;
use std::cell::RefCell;
use std::marker::PhantomData;

use super::super::super::api::*;
use super::*;
use implementation_support::common::{split_off_greater_set};

#[derive (Copy, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, Debug)]
pub struct GetVarying;
impl StewardData for GetVarying {}

#[derive (Clone, Serialize, Deserialize, Debug)]
pub struct SimpleTimeline <VaryingData: StewardData, B: Basics> {
  // Hacky workaround for https://github.com/rust-lang/rust/issues/41617 (see https://github.com/serde-rs/serde/issues/943)
  #[serde(deserialize_with = "::serde::Deserialize::deserialize")]
  changes: Vec<(EventHandle <B>, Option <VaryingData>)>,
  #[serde(deserialize_with = "::serde::Deserialize::deserialize")]
  other_dependent_events: RefCell<BTreeSet<EventHandle<B>>>,
}

impl <VaryingData: StewardData, B: Basics> SimpleTimeline <VaryingData, B> {
  pub fn new ()->Self {
    SimpleTimeline {
      changes: Vec::new(),
      other_dependent_events: RefCell::new (BTreeSet::new()),
    }
  }

  fn invalidate_after <Steward: TimeSteward <Basics = B, EventHandle = EventHandle<B>>, Accessor: InvalidationAccessor<Steward = Steward>> (&self, time: &ExtendedTime <B>, accessor: & Accessor) {
    let mut dependencies = self.other_dependent_events.borrow_mut();
    let removed = split_off_greater_set (&mut *dependencies, time);
    for event in removed {
      accessor.invalidate(&event);
    }
    for change in self.changes.iter().rev() {
      let event = &change.0;
      if event.extended_time() <= time {
        break
      }
      accessor.invalidate(&event);
    }
  }

  fn remove_from (&mut self, time: &ExtendedTime <B>) {
    while let Some (change) = self.changes.pop() {
      if change.0.extended_time() < time {
        self.changes.push (change);
        break
      }
    }
  }
}

impl <VaryingData: StewardData, B: Basics> DataTimeline for SimpleTimeline <VaryingData, B> {
  type Basics = B;

  fn clone_for_snapshot (&self, time: &ExtendedTime <Self::Basics>)->Self {
    let slice = match self.changes.binary_search_by_key (&time, | change | change.0.extended_time()) {
      Ok(index) => &self.changes [index.saturating_sub (1).. index+1],
      Err (index) => &self.changes [index.saturating_sub (1)..index],
    };
    SimpleTimeline {
      changes: slice.to_vec(),
      other_dependent_events: RefCell::new (BTreeSet::new()),
    }
  }

  fn forget_before (&mut self, time: &ExtendedTime <Self::Basics>) {
    let mut dependencies = self.other_dependent_events.borrow_mut();
    let retained = dependencies.split_off (time);
    mem::replace (&mut*dependencies, retained);

    if self.changes.len() > 0 && self.changes [self.changes.len()/2].0.extended_time() < time {
      let keep_from_index = match self.changes.binary_search_by_key (&time, | change | change.0.extended_time()) {
        Ok (index) => index,
        Err(index) => index - 1,
      };
      self.changes = self.changes.split_off (keep_from_index);
    }
  }
}
impl <VaryingData: StewardData, B: Basics> DataTimelineQueriableWith<GetVarying> for SimpleTimeline <VaryingData, B> {
  type QueryResult = Option <(ExtendedTime <B>, VaryingData)>;

  fn query (&self, _: &GetVarying, time: &ExtendedTime <Self::Basics>, offset: QueryOffset)->Self::QueryResult {
    let previous_change_index = match self.changes.binary_search_by_key (&time, | change | change.0.extended_time()) {
      Ok(index) => match offset {QueryOffset::After => index, QueryOffset::Before => index.wrapping_sub (1)},
      Err (index) => index.wrapping_sub (1),
    };
    self.changes.get (previous_change_index).and_then (| change | change.1.as_ref().map (| data | (change.0.extended_time().clone(), data.clone())))
  }
}

pub fn tracking_query <B: Basics, VaryingData: StewardData, Steward: TimeSteward <Basics = B, EventHandle = EventHandle<B>>, Accessor: EventAccessor <Steward = Steward>> (accessor: & Accessor, handle: & DataTimelineCell <SimpleTimeline <VaryingData, Steward::Basics>>, offset: QueryOffset)->Option <(ExtendedTime <Steward::Basics>, VaryingData)> {
  accessor.modify (handle, move |timeline| {
    let mut dependencies = timeline.other_dependent_events.borrow_mut();
    dependencies.insert (accessor.handle().clone());
  });
  accessor.query (handle, &GetVarying, offset)
}
pub fn modify_simple_timeline <B: Basics, VaryingData: StewardData, Steward: TimeSteward <Basics = B, EventHandle = EventHandle<B>>, Accessor: EventAccessor <Steward = Steward>> (accessor: & Accessor, handle: & DataTimelineCell <SimpleTimeline <VaryingData, Steward::Basics>>, modification: Option <VaryingData>) {
  match accessor.query (handle, &GetVarying, QueryOffset::After) {
    Some((time, data)) =>
      if let Some (new_data) = modification.as_ref() {
        if &time == accessor.extended_now() && &data == new_data {
          return
        }
      },
    None =>
      if modification.is_none() {return}
  };
  accessor.invalidate (| invalidator | {
    invalidator.peek(handle).invalidate_after (accessor.extended_now(), invalidator);
  });
  accessor.modify (handle, move |timeline| {
    timeline.remove_from (accessor.extended_now());
    timeline.changes.push ((accessor.handle().clone(), modification));
  });
}
pub fn unmodify_simple_timeline <B: Basics, VaryingData: StewardData, Steward: TimeSteward <Basics = B, EventHandle = EventHandle<B>>, Accessor: EventAccessor <Steward = Steward>> (accessor: & Accessor, handle: & DataTimelineCell <SimpleTimeline <VaryingData, Steward::Basics>>) {
  if let Some((time, _)) = accessor.query (handle, &GetVarying, QueryOffset::After) { if &time == accessor.extended_now() {
    accessor.invalidate (| invalidator | {
      invalidator.peek(handle).invalidate_after (accessor.extended_now(), invalidator);
    });
    accessor.modify (handle, move |timeline| {
      timeline.remove_from (accessor.extended_now());
    });
  }}
}

} //mod

  };
}

<commit_msg>Optimize SimpleTimeline a bit for the most common case<commit_after>#[doc (hidden)]
#[macro_export]
macro_rules! time_steward_define_simple_timeline {
  () => {
pub mod simple_timeline {
use std::collections::BTreeSet;
use std::mem;
use std::cell::RefCell;
use std::marker::PhantomData;

use super::super::super::api::*;
use super::*;
use implementation_support::common::{split_off_greater_set};

#[derive (Copy, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, Debug)]
pub struct GetVarying;
impl StewardData for GetVarying {}

#[derive (Clone, Serialize, Deserialize, Debug)]
pub struct SimpleTimeline <VaryingData: StewardData, B: Basics> {
  // Hacky workaround for https://github.com/rust-lang/rust/issues/41617 (see https://github.com/serde-rs/serde/issues/943)
  #[serde(deserialize_with = "::serde::Deserialize::deserialize")]
  changes: Vec<(EventHandle <B>, Option <VaryingData>)>,
  #[serde(deserialize_with = "::serde::Deserialize::deserialize")]
  other_dependent_events: RefCell<BTreeSet<EventHandle<B>>>,
}

impl <VaryingData: StewardData, B: Basics> SimpleTimeline <VaryingData, B> {
  pub fn new ()->Self {
    SimpleTimeline {
      changes: Vec::new(),
      other_dependent_events: RefCell::new (BTreeSet::new()),
    }
  }

  fn invalidate_after <Steward: TimeSteward <Basics = B, EventHandle = EventHandle<B>>, Accessor: InvalidationAccessor<Steward = Steward>> (&self, time: &ExtendedTime <B>, accessor: & Accessor) {
    let mut dependencies = self.other_dependent_events.borrow_mut();
    let removed = split_off_greater_set (&mut *dependencies, time);
    for event in removed {
      accessor.invalidate(&event);
    }
    for change in self.changes.iter().rev() {
      let event = &change.0;
      if event.extended_time() <= time {
        break
      }
      accessor.invalidate(&event);
    }
  }

  fn remove_from (&mut self, time: &ExtendedTime <B>) {
    while let Some (change) = self.changes.pop() {
      if change.0.extended_time() < time {
        self.changes.push (change);
        break
      }
    }
  }

  fn search_changes (&self, time: &ExtendedTime <B>) -> Result <usize, usize> {
    // search at the end first, because we are usually in the present.
    match self.changes.last() {
      None => return Err(0),
      Some (&(ref event, ref data)) => {
        match event.extended_time().cmp(time) {
          Ordering::Greater => (),
          Ordering::Equal => return Ok(self.changes.len()-1),
          Ordering::Less => return Err(self.changes.len()),
        }
      },
    }
    self.changes[..self.changes.len()-1].binary_search_by_key (&time, | change | change.0.extended_time())
  }
}

impl <VaryingData: StewardData, B: Basics> DataTimeline for SimpleTimeline <VaryingData, B> {
  type Basics = B;

  fn clone_for_snapshot (&self, time: &ExtendedTime <Self::Basics>)->Self {
    let slice = match self.search_changes(&time) {
      Ok(index) => &self.changes [index.saturating_sub (1).. index+1],
      Err (index) => &self.changes [index.saturating_sub (1)..index],
    };
    SimpleTimeline {
      changes: slice.to_vec(),
      other_dependent_events: RefCell::new (BTreeSet::new()),
    }
  }

  fn forget_before (&mut self, time: &ExtendedTime <Self::Basics>) {
    let mut dependencies = self.other_dependent_events.borrow_mut();
    let retained = dependencies.split_off (time);
    mem::replace (&mut*dependencies, retained);

    if self.changes.len() > 0 && self.changes [self.changes.len()/2].0.extended_time() < time {
      let keep_from_index = match self.search_changes(&time) {
        Ok (index) => index,
        Err(index) => index - 1,
      };
      self.changes = self.changes.split_off (keep_from_index);
    }
  }
}
impl <VaryingData: StewardData, B: Basics> DataTimelineQueriableWith<GetVarying> for SimpleTimeline <VaryingData, B> {
  type QueryResult = Option <(ExtendedTime <B>, VaryingData)>;

  fn query (&self, _: &GetVarying, time: &ExtendedTime <Self::Basics>, offset: QueryOffset)->Self::QueryResult {
    let previous_change_index = match self.search_changes(&time) {
      Ok(index) => match offset {QueryOffset::After => index, QueryOffset::Before => index.wrapping_sub (1)},
      Err (index) => index.wrapping_sub (1),
    };
    self.changes.get (previous_change_index).and_then (| change | change.1.as_ref().map (| data | (change.0.extended_time().clone(), data.clone())))
  }
}

pub fn tracking_query <B: Basics, VaryingData: StewardData, Steward: TimeSteward <Basics = B, EventHandle = EventHandle<B>>, Accessor: EventAccessor <Steward = Steward>> (accessor: & Accessor, handle: & DataTimelineCell <SimpleTimeline <VaryingData, Steward::Basics>>, offset: QueryOffset)->Option <(ExtendedTime <Steward::Basics>, VaryingData)> {
  accessor.modify (handle, move |timeline| {
    let mut dependencies = timeline.other_dependent_events.borrow_mut();
    dependencies.insert (accessor.handle().clone());
  });
  accessor.query (handle, &GetVarying, offset)
}
pub fn modify_simple_timeline <B: Basics, VaryingData: StewardData, Steward: TimeSteward <Basics = B, EventHandle = EventHandle<B>>, Accessor: EventAccessor <Steward = Steward>> (accessor: & Accessor, handle: & DataTimelineCell <SimpleTimeline <VaryingData, Steward::Basics>>, modification: Option <VaryingData>) {
  match accessor.query (handle, &GetVarying, QueryOffset::After) {
    Some((time, data)) =>
      if let Some (new_data) = modification.as_ref() {
        if &time == accessor.extended_now() && &data == new_data {
          return
        }
      },
    None =>
      if modification.is_none() {return}
  };
  accessor.invalidate (| invalidator | {
    invalidator.peek(handle).invalidate_after (accessor.extended_now(), invalidator);
  });
  accessor.modify (handle, move |timeline| {
    timeline.remove_from (accessor.extended_now());
    timeline.changes.push ((accessor.handle().clone(), modification));
  });
}
pub fn unmodify_simple_timeline <B: Basics, VaryingData: StewardData, Steward: TimeSteward <Basics = B, EventHandle = EventHandle<B>>, Accessor: EventAccessor <Steward = Steward>> (accessor: & Accessor, handle: & DataTimelineCell <SimpleTimeline <VaryingData, Steward::Basics>>) {
  if let Some((time, _)) = accessor.query (handle, &GetVarying, QueryOffset::After) { if &time == accessor.extended_now() {
    accessor.invalidate (| invalidator | {
      invalidator.peek(handle).invalidate_after (accessor.extended_now(), invalidator);
    });
    accessor.modify (handle, move |timeline| {
      timeline.remove_from (accessor.extended_now());
    });
  }}
}

} //mod

  };
}

<|endoftext|>


<commit_before>use std::cmp::min;
use std::collections::{BTreeMap, HashSet};
use std::convert::TryInto;
use std::fmt;
use std::sync::Arc;
use std::time::Duration;

use bazel_protos::{self, call_option};
use bytes::{Bytes, BytesMut};
use futures::compat::Future01CompatExt;
use futures::future::{FutureExt, TryFutureExt};
use futures01::{future, Future, Sink, Stream};
use hashing::Digest;
use log::Level;
use serverset::{retry, Serverset};
use workunit_store::with_workunit;

use super::{BackoffConfig, EntryType};

#[derive(Clone)]
pub struct ByteStore {
  instance_name: Option<String>,
  chunk_size_bytes: usize,
  upload_timeout: Duration,
  rpc_attempts: usize,
  env: Arc<grpcio::Environment>,
  serverset: Serverset<grpcio::Channel>,
  headers: BTreeMap<String, String>,
}

impl fmt::Debug for ByteStore {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    write!(f, "ByteStore(name={:?})", self.instance_name)
  }
}

impl ByteStore {
  pub fn new(
    cas_addresses: Vec<String>,
    instance_name: Option<String>,
    root_ca_certs: Option<Vec<u8>>,
    oauth_bearer_token: Option<String>,
    thread_count: usize,
    chunk_size_bytes: usize,
    upload_timeout: Duration,
    backoff_config: BackoffConfig,
    rpc_retries: usize,
    connection_limit: usize,
  ) -> Result<ByteStore, String> {
    let env = Arc::new(grpcio::Environment::new(thread_count));
    let env2 = env.clone();

    let connect = move |cas_address: &str| {
      let builder = grpcio::ChannelBuilder::new(env2.clone());
      if let Some(ref root_ca_certs) = root_ca_certs {
        let creds = grpcio::ChannelCredentialsBuilder::new()
          .root_cert(root_ca_certs.clone())
          .build();
        builder.secure_connect(cas_address, creds)
      } else {
        builder.connect(cas_address)
      }
    };

    let serverset = Serverset::new(cas_addresses, connect, connection_limit, backoff_config)?;

    Ok(ByteStore {
      instance_name,
      chunk_size_bytes,
      upload_timeout,
      rpc_attempts: rpc_retries + 1,
      env,
      serverset,
      headers: oauth_bearer_token
        .iter()
        .map(|t| (String::from("authorization"), format!("Bearer {}", t)))
        .collect(),
    })
  }

  async fn with_byte_stream_client<
    Value: Send,
    Fut: std::future::Future<Output = Result<Value, String>>,
    F: Fn(bazel_protos::bytestream_grpc::ByteStreamClient) -> Fut,
  >(
    &self,
    f: F,
  ) -> Result<Value, String> {
    retry::all_errors_immediately(&self.serverset, self.rpc_attempts, move |channel| {
      f(bazel_protos::bytestream_grpc::ByteStreamClient::new(
        channel,
      ))
    })
    .await
  }

  async fn with_cas_client<
    Value: Send,
    Fut: std::future::Future<Output = Result<Value, String>>,
    F: Fn(bazel_protos::remote_execution_grpc::ContentAddressableStorageClient) -> Fut,
  >(
    &self,
    f: F,
  ) -> Result<Value, String> {
    retry::all_errors_immediately(&self.serverset, self.rpc_attempts, move |channel| {
      f(bazel_protos::remote_execution_grpc::ContentAddressableStorageClient::new(channel))
    })
    .await
  }

  pub async fn store_bytes(&self, bytes: &[u8]) -> Result<Digest, String> {
    let len = bytes.len();
    let digest = Digest::of_bytes(&bytes);
    let resource_name = format!(
      "{}/uploads/{}/blobs/{}/{}",
      self.instance_name.clone().unwrap_or_default(),
      uuid::Uuid::new_v4(),
      digest.0,
      digest.1,
    );
    let workunit_name = format!("store_bytes({})", resource_name.clone());
    let metadata = workunit_store::WorkunitMetadata::with_level(Level::Debug);
    let store = self.clone();
    let result_future = self.with_byte_stream_client(move |client| {
      let resource_name = resource_name.clone();
      let store = store.clone();
      async move {
        let (sender, receiver) = client
          .write_opt(call_option(&store.headers, None)?.timeout(store.upload_timeout))
          .map_err(|err| {
            format!(
              "Error attempting to connect to upload digest {:?}: {:?}",
              digest, err
            )
          })?;

        let chunk_size_bytes = store.chunk_size_bytes;
        let resource_name = resource_name.clone();
        let stream = futures01::stream::unfold::<_, _, future::FutureResult<_, grpcio::Error>, _>(
          (0, false),
          move |(offset, has_sent_any)| {
            if offset >= bytes.len() && has_sent_any {
              None
            } else {
              let mut req = bazel_protos::bytestream::WriteRequest::new();
              req.set_resource_name(resource_name.clone());
              req.set_write_offset(offset as i64);
              let next_offset = min(offset + chunk_size_bytes, bytes.len());
              req.set_finish_write(next_offset == bytes.len());
              req.set_data(Bytes::from(&bytes[offset..next_offset]));
              Some(future::ok((
                (req, grpcio::WriteFlags::default()),
                (next_offset, true),
              )))
            }
          },
        );

        sender
          .send_all(stream)
          .map(|_| ())
          .or_else(move |e| {
            match e {
              // Some implementations of the remote execution API early-return if the blob has
              // been concurrently uploaded by another client. In this case, they return a
              // WriteResponse with a committed_size equal to the digest's entire size before
              // closing the stream.
              // Because the server then closes the stream, the client gets an RpcFinished
              // error in this case. We ignore this, and will later on verify that the
              // committed_size we received from the server is equal to the expected one. If
              // these are not equal, the upload will be considered a failure at that point.
              // Whether this type of response will become part of the official API is up for
              // discussion: see
              // https://groups.google.com/d/topic/remote-execution-apis/NXUe3ItCw68/discussion.
              grpcio::Error::RpcFinished(None) => Ok(()),
              e => Err(format!(
                "Error attempting to upload digest {:?}: {:?}",
                digest, e
              )),
            }
          })
          .compat()
          .await?;

        let received = receiver
          .map_err(move |e| {
            format!(
              "Error from server when uploading digest {:?}: {:?}",
              digest, e
            )
          })
          .compat()
          .await?;

        if received.get_committed_size() == len as i64 {
          Ok(digest)
        } else {
          Err(format!(
            "Uploading file with digest {:?}: want committed size {} but got {}",
            digest,
            len,
            received.get_committed_size()
          ))
        }
      }
    });

    if let Some(workunit_state) = workunit_store::get_workunit_state() {
      let store = workunit_state.store;
      with_workunit(store, workunit_name, metadata, result_future, |_, md| md).await
    } else {
      result_future.await
    }
  }

  pub async fn load_bytes_with<
    T: Send + 'static,
    F: Fn(Bytes) -> T + Send + Sync + Clone + 'static,
  >(
    &self,
    _entry_type: EntryType,
    digest: Digest,
    f: F,
  ) -> Result<Option<T>, String> {
    let store = self.clone();
    let resource_name = format!(
      "{}/blobs/{}/{}",
      store.instance_name.clone().unwrap_or_default(),
      digest.0,
      digest.1
    );
    let workunit_name = format!("load_bytes_with({})", resource_name.clone());
    let metadata = workunit_store::WorkunitMetadata::with_level(Level::Debug);
    let result_future = self.with_byte_stream_client(move |client| {
      let resource_name = resource_name.clone();
      let store = store.clone();
      let f = f.clone();
      async move {
        let stream = client
          .read_opt(
            &{
              let mut req = bazel_protos::bytestream::ReadRequest::new();
              req.set_resource_name(resource_name.clone());
              req.set_read_offset(0);
              // 0 means no limit.
              req.set_read_limit(0);
              req
            },
            call_option(&store.headers, None)?,
          )
          .map_err(|err| format!("Error making CAS read request for {:?}: {:?}", digest, err))?;

        let bytes_res = stream
          .fold(BytesMut::with_capacity(digest.1), move |mut bytes, r| {
            bytes.extend_from_slice(&r.data);
            future::ok::<_, grpcio::Error>(bytes)
          })
          .compat()
          .await;

        // We ensure that we hold onto the client until after we've consumed the stream as a
        // workaround for https://github.com/pingcap/grpc-rs/issues/123
        std::mem::drop(client);

        let maybe_bytes = match bytes_res {
          Ok(bytes) => Some(bytes.freeze()),
          Err(grpcio::Error::RpcFailure(grpcio::RpcStatus {
            status: grpcio::RpcStatusCode::NOT_FOUND,
            ..
          })) => None,
          Err(e) => {
            return Err(format!(
              "Error from server in response to CAS read request: {:?}",
              e
            ))
          }
        };

        Ok(maybe_bytes.map(f))
      }
    });

    if let Some(workunit_state) = workunit_store::get_workunit_state() {
      let store = workunit_state.store;
      with_workunit(store, workunit_name, metadata, result_future, |_, md| md).await
    } else {
      result_future.await
    }
  }

  ///
  /// Given a collection of Digests (digests),
  /// returns the set of digests from that collection not present in the CAS.
  ///
  pub fn list_missing_digests(
    &self,
    request: bazel_protos::remote_execution::FindMissingBlobsRequest,
  ) -> impl Future<Item = HashSet<Digest>, Error = String> {
    let store = self.clone();
    let workunit_name = format!(
      "list_missing_digests({})",
      store.instance_name.clone().unwrap_or_default()
    );
    let metadata = workunit_store::WorkunitMetadata::with_level(Level::Debug);
    let result_future = async move {
      let store2 = store.clone();
      store2
        .with_cas_client(move |client| {
          let request = request.clone();
          let store = store.clone();
          async move {
            let response = client
              .find_missing_blobs_opt(&request, call_option(&store.headers, None)?)
              .map_err(|err| {
                format!(
                  "Error from server in response to find_missing_blobs_request: {:?}",
                  err
                )
              })?;
            response
              .get_missing_blob_digests()
              .iter()
              .map(|digest| digest.try_into())
              .collect::<Result<HashSet<_>, _>>()
          }
        })
        .await
    };
    async {
      if let Some(workunit_state) = workunit_store::get_workunit_state() {
        let store = workunit_state.store;
        with_workunit(store, workunit_name, metadata, result_future, |_, md| md).await
      } else {
        result_future.await
      }
    }
    .boxed()
    .compat()
  }

  pub(super) fn find_missing_blobs_request<'a, Digests: Iterator<Item = &'a Digest>>(
    &self,
    digests: Digests,
  ) -> bazel_protos::remote_execution::FindMissingBlobsRequest {
    let mut request = bazel_protos::remote_execution::FindMissingBlobsRequest::new();
    if let Some(ref instance_name) = self.instance_name {
      request.set_instance_name(instance_name.clone());
    }
    for digest in digests {
      request.mut_blob_digests().push(digest.into());
    }
    request
  }
}
<commit_msg>close write stream for CAS writes via ByteStream API (#10395)<commit_after>use std::cmp::min;
use std::collections::{BTreeMap, HashSet};
use std::convert::TryInto;
use std::fmt;
use std::sync::Arc;
use std::time::Duration;

use bazel_protos::{self, call_option};
use bytes::{Bytes, BytesMut};
use futures::compat::{Future01CompatExt, Sink01CompatExt};
use futures::future::{FutureExt, TryFutureExt};
use futures::sink::SinkExt;
use futures01::{future, Future, Stream};
use hashing::Digest;
use log::Level;
use serverset::{retry, Serverset};
use workunit_store::with_workunit;

use super::{BackoffConfig, EntryType};

#[derive(Clone)]
pub struct ByteStore {
  instance_name: Option<String>,
  chunk_size_bytes: usize,
  upload_timeout: Duration,
  rpc_attempts: usize,
  env: Arc<grpcio::Environment>,
  serverset: Serverset<grpcio::Channel>,
  headers: BTreeMap<String, String>,
}

impl fmt::Debug for ByteStore {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    write!(f, "ByteStore(name={:?})", self.instance_name)
  }
}

impl ByteStore {
  pub fn new(
    cas_addresses: Vec<String>,
    instance_name: Option<String>,
    root_ca_certs: Option<Vec<u8>>,
    oauth_bearer_token: Option<String>,
    thread_count: usize,
    chunk_size_bytes: usize,
    upload_timeout: Duration,
    backoff_config: BackoffConfig,
    rpc_retries: usize,
    connection_limit: usize,
  ) -> Result<ByteStore, String> {
    let env = Arc::new(grpcio::Environment::new(thread_count));
    let env2 = env.clone();

    let connect = move |cas_address: &str| {
      let builder = grpcio::ChannelBuilder::new(env2.clone());
      if let Some(ref root_ca_certs) = root_ca_certs {
        let creds = grpcio::ChannelCredentialsBuilder::new()
          .root_cert(root_ca_certs.clone())
          .build();
        builder.secure_connect(cas_address, creds)
      } else {
        builder.connect(cas_address)
      }
    };

    let serverset = Serverset::new(cas_addresses, connect, connection_limit, backoff_config)?;

    Ok(ByteStore {
      instance_name,
      chunk_size_bytes,
      upload_timeout,
      rpc_attempts: rpc_retries + 1,
      env,
      serverset,
      headers: oauth_bearer_token
        .iter()
        .map(|t| (String::from("authorization"), format!("Bearer {}", t)))
        .collect(),
    })
  }

  async fn with_byte_stream_client<
    Value: Send,
    Fut: std::future::Future<Output = Result<Value, String>>,
    F: Fn(bazel_protos::bytestream_grpc::ByteStreamClient) -> Fut,
  >(
    &self,
    f: F,
  ) -> Result<Value, String> {
    retry::all_errors_immediately(&self.serverset, self.rpc_attempts, move |channel| {
      f(bazel_protos::bytestream_grpc::ByteStreamClient::new(
        channel,
      ))
    })
    .await
  }

  async fn with_cas_client<
    Value: Send,
    Fut: std::future::Future<Output = Result<Value, String>>,
    F: Fn(bazel_protos::remote_execution_grpc::ContentAddressableStorageClient) -> Fut,
  >(
    &self,
    f: F,
  ) -> Result<Value, String> {
    retry::all_errors_immediately(&self.serverset, self.rpc_attempts, move |channel| {
      f(bazel_protos::remote_execution_grpc::ContentAddressableStorageClient::new(channel))
    })
    .await
  }

  pub async fn store_bytes(&self, bytes: &[u8]) -> Result<Digest, String> {
    let len = bytes.len();
    let digest = Digest::of_bytes(&bytes);
    let resource_name = format!(
      "{}/uploads/{}/blobs/{}/{}",
      self.instance_name.clone().unwrap_or_default(),
      uuid::Uuid::new_v4(),
      digest.0,
      digest.1,
    );
    let workunit_name = format!("store_bytes({})", resource_name.clone());
    let metadata = workunit_store::WorkunitMetadata::with_level(Level::Debug);
    let store = self.clone();
    let result_future = self.with_byte_stream_client(move |client| {
      let resource_name = resource_name.clone();
      let store = store.clone();
      async move {
        let (sender, receiver) = client
          .write_opt(call_option(&store.headers, None)?.timeout(store.upload_timeout))
          .map_err(|err| {
            format!(
              "Error attempting to connect to upload digest {:?}: {:?}",
              digest, err
            )
          })?;

        let mut sender = sender.sink_compat();

        let chunk_size_bytes = store.chunk_size_bytes;
        let resource_name = resource_name.clone();
        let mut stream = futures::stream::unfold((0, false), move |(offset, has_sent_any)| {
          if offset >= bytes.len() && has_sent_any {
            futures::future::ready(None)
          } else {
            let mut req = bazel_protos::bytestream::WriteRequest::new();
            req.set_resource_name(resource_name.clone());
            req.set_write_offset(offset as i64);
            let next_offset = min(offset + chunk_size_bytes, bytes.len());
            req.set_finish_write(next_offset == bytes.len());
            req.set_data(Bytes::from(&bytes[offset..next_offset]));
            futures::future::ready(Some((
              Ok((req, grpcio::WriteFlags::default())),
              (next_offset, true),
            )))
          }
        });

        sender.send_all(&mut stream).await.or_else(move |e| {
          match e {
            // Some implementations of the remote execution API early-return if the blob has
            // been concurrently uploaded by another client. In this case, they return a
            // WriteResponse with a committed_size equal to the digest's entire size before
            // closing the stream.
            // Because the server then closes the stream, the client gets an RpcFinished
            // error in this case. We ignore this, and will later on verify that the
            // committed_size we received from the server is equal to the expected one. If
            // these are not equal, the upload will be considered a failure at that point.
            // Whether this type of response will become part of the official API is up for
            // discussion: see
            // https://groups.google.com/d/topic/remote-execution-apis/NXUe3ItCw68/discussion.
            grpcio::Error::RpcFinished(None) => Ok(()),
            e => Err(format!(
              "Error attempting to upload digest {:?}: {:?}",
              digest, e
            )),
          }
        })?;

        // The gRPC library cancels streams on drop; closes must be explicit. Not closing
        // the stream caused the BuildGrid CAS server to generate errors on writes
        // when the stream was cancelled because it was not closed explicitly.
        sender.close().await.or_else(|err| {
          match err {
            // Some servers (e.g., RBE) may have already closed the stream for the early
            // return reason identified previously. Treat this condition as a successful close.
            grpcio::Error::RpcFinished(None) => Ok(()),
            e => Err(format!(
              "Error from server when uploading digest {:?}: {:?}",
              digest, e
            )),
          }
        })?;

        let received = receiver
          .map_err(move |e| {
            format!(
              "Error from server when uploading digest {:?}: {:?}",
              digest, e
            )
          })
          .compat()
          .await?;

        if received.get_committed_size() == len as i64 {
          Ok(digest)
        } else {
          Err(format!(
            "Uploading file with digest {:?}: want committed size {} but got {}",
            digest,
            len,
            received.get_committed_size()
          ))
        }
      }
    });

    if let Some(workunit_state) = workunit_store::get_workunit_state() {
      let store = workunit_state.store;
      with_workunit(store, workunit_name, metadata, result_future, |_, md| md).await
    } else {
      result_future.await
    }
  }

  pub async fn load_bytes_with<
    T: Send + 'static,
    F: Fn(Bytes) -> T + Send + Sync + Clone + 'static,
  >(
    &self,
    _entry_type: EntryType,
    digest: Digest,
    f: F,
  ) -> Result<Option<T>, String> {
    let store = self.clone();
    let resource_name = format!(
      "{}/blobs/{}/{}",
      store.instance_name.clone().unwrap_or_default(),
      digest.0,
      digest.1
    );
    let workunit_name = format!("load_bytes_with({})", resource_name.clone());
    let metadata = workunit_store::WorkunitMetadata::with_level(Level::Debug);
    let result_future = self.with_byte_stream_client(move |client| {
      let resource_name = resource_name.clone();
      let store = store.clone();
      let f = f.clone();
      async move {
        let stream = client
          .read_opt(
            &{
              let mut req = bazel_protos::bytestream::ReadRequest::new();
              req.set_resource_name(resource_name.clone());
              req.set_read_offset(0);
              // 0 means no limit.
              req.set_read_limit(0);
              req
            },
            call_option(&store.headers, None)?,
          )
          .map_err(|err| format!("Error making CAS read request for {:?}: {:?}", digest, err))?;

        let bytes_res = stream
          .fold(BytesMut::with_capacity(digest.1), move |mut bytes, r| {
            bytes.extend_from_slice(&r.data);
            future::ok::<_, grpcio::Error>(bytes)
          })
          .compat()
          .await;

        // We ensure that we hold onto the client until after we've consumed the stream as a
        // workaround for https://github.com/pingcap/grpc-rs/issues/123
        std::mem::drop(client);

        let maybe_bytes = match bytes_res {
          Ok(bytes) => Some(bytes.freeze()),
          Err(grpcio::Error::RpcFailure(grpcio::RpcStatus {
            status: grpcio::RpcStatusCode::NOT_FOUND,
            ..
          })) => None,
          Err(e) => {
            return Err(format!(
              "Error from server in response to CAS read request: {:?}",
              e
            ))
          }
        };

        Ok(maybe_bytes.map(f))
      }
    });

    if let Some(workunit_state) = workunit_store::get_workunit_state() {
      let store = workunit_state.store;
      with_workunit(store, workunit_name, metadata, result_future, |_, md| md).await
    } else {
      result_future.await
    }
  }

  ///
  /// Given a collection of Digests (digests),
  /// returns the set of digests from that collection not present in the CAS.
  ///
  pub fn list_missing_digests(
    &self,
    request: bazel_protos::remote_execution::FindMissingBlobsRequest,
  ) -> impl Future<Item = HashSet<Digest>, Error = String> {
    let store = self.clone();
    let workunit_name = format!(
      "list_missing_digests({})",
      store.instance_name.clone().unwrap_or_default()
    );
    let metadata = workunit_store::WorkunitMetadata::with_level(Level::Debug);
    let result_future = async move {
      let store2 = store.clone();
      store2
        .with_cas_client(move |client| {
          let request = request.clone();
          let store = store.clone();
          async move {
            let response = client
              .find_missing_blobs_opt(&request, call_option(&store.headers, None)?)
              .map_err(|err| {
                format!(
                  "Error from server in response to find_missing_blobs_request: {:?}",
                  err
                )
              })?;
            response
              .get_missing_blob_digests()
              .iter()
              .map(|digest| digest.try_into())
              .collect::<Result<HashSet<_>, _>>()
          }
        })
        .await
    };
    async {
      if let Some(workunit_state) = workunit_store::get_workunit_state() {
        let store = workunit_state.store;
        with_workunit(store, workunit_name, metadata, result_future, |_, md| md).await
      } else {
        result_future.await
      }
    }
    .boxed()
    .compat()
  }

  pub(super) fn find_missing_blobs_request<'a, Digests: Iterator<Item = &'a Digest>>(
    &self,
    digests: Digests,
  ) -> bazel_protos::remote_execution::FindMissingBlobsRequest {
    let mut request = bazel_protos::remote_execution::FindMissingBlobsRequest::new();
    if let Some(ref instance_name) = self.instance_name {
      request.set_instance_name(instance_name.clone());
    }
    for digest in digests {
      request.mut_blob_digests().push(digest.into());
    }
    request
  }
}
<|endoftext|>


<commit_before>export Msg, Prefetch, GetImage, Exit;
export ImageResponseMsg, ImageReady, ImageNotReady;
export ImageCacheTask;
export image_cache_task;

import image::base::{Image, load_from_memory, test_image_bin};
import std::net::url::url;
import util::url::{make_url, UrlMap, url_map};
import comm::{chan, port};
import task::{spawn, spawn_listener};
import resource::resource_task;
import resource_task::ResourceTask;
import core::arc::arc;
import clone_arc = core::arc::clone;

enum Msg {
    Prefetch(url),
    GetImage(url, chan<ImageResponseMsg>),
    StoreImage(url, arc<~Image>),
    Exit
}

enum ImageResponseMsg {
    ImageReady(arc<~Image>),
    ImageNotReady
}

type ImageCacheTask = chan<Msg>;

fn image_cache_task(resource_task: ResourceTask) -> ImageCacheTask {
    do spawn_listener |from_client| {
        ImageCache {
            resource_task: resource_task,
            from_client: from_client,
            prefetch_map: url_map(),
            future_image_map: url_map(),
            image_map: url_map()
        }.run();
    }
}

struct ImageCache {
    resource_task: ResourceTask;
    from_client: port<Msg>;
    prefetch_map: UrlMap<@PrefetchData>;
    future_image_map: UrlMap<@FutureData>;
    image_map: UrlMap<@arc<~Image>>;
}

struct PrefetchData {
    response_port: port<resource_task::ProgressMsg>;
    mut data: ~[u8];
}

struct FutureData {
    mut waiters: ~[chan<ImageResponseMsg>];
}

impl ImageCache {

    fn run() {

        loop {
            match self.from_client.recv() {
              Prefetch(url) => self.prefetch(url),
              GetImage(url, response) => self.get_image(url, response),
              StoreImage(url, image) => self.store_image(url, &image),
              Exit => break
            }
        }
    }

    /*priv*/ fn prefetch(url: url) {
        if self.image_map.contains_key(url) {
            // We've already decoded this image
            return
        }

        if self.future_image_map.contains_key(url) {
            // We've already begun decoding this image
            return
        }

        if self.prefetch_map.contains_key(url) {
            // We're already waiting for this image
            return
        }

        let response_port = port();
        self.resource_task.send(resource_task::Load(url, response_port.chan()));

        let prefetch_data = @PrefetchData {
            response_port: response_port,
            data: ~[]
        };

        self.prefetch_map.insert(url, prefetch_data);
    }

    /*priv*/ fn get_image(url: url, response: chan<ImageResponseMsg>) {
        match self.image_map.find(url) {
          some(image) => {
            response.send(ImageReady(clone_arc(image)));
            return
          }
          none => ()
        }

        match self.future_image_map.find(url) {
          some(future_data) => {
            // We've started decoding this image but haven't recieved it back yet.
            // Put this client on the wait list
            vec::push(future_data.waiters, response);
            return
          }
          none => ()
        }

        match self.prefetch_map.find(url) {
          some(prefetch_data) => {

            let mut image_sent = false;

            while prefetch_data.response_port.peek() {
                match prefetch_data.response_port.recv() {
                  resource_task::Payload(data) => {
                    prefetch_data.data += data;
                  }
                  resource_task::Done(result::ok(*)) => {
                    // We've got the entire image binary
                    let mut data = ~[];
                    data <-> prefetch_data.data;
                    let data <- data; // freeze for capture

                    let to_cache = self.from_client.chan();

                    do spawn {
                        let image = arc(~load_from_memory(data));
                        // Send the image to the original requester
                        response.send(ImageReady(clone_arc(&image)));
                        to_cache.send(StoreImage(url, clone_arc(&image)));
                    }

                    let future_data = @FutureData {
                        waiters: ~[]
                    };

                    self.prefetch_map.remove(url);
                    self.future_image_map.insert(url, future_data);

                    image_sent = true;
                    break;
                  }
                  resource_task::Done(result::err(*)) => {
                    // FIXME: need to actually report the failed image load
                    self.prefetch_map.remove(url);
                    break;
                  }
                }
            }

            if !image_sent {
                response.send(ImageNotReady);
            }
          }
          none => fail ~"got a request for image data without prefetch"
        }
    }

    /*priv*/ fn store_image(url: url, image: &arc<~Image>) {
        match self.future_image_map.find(url) {
          some(future_data) => {

            let mut waiters = ~[];
            waiters <-> future_data.waiters;

            // Send the image to all those who requested it while
            // it was being decoded
            for waiters.each |waiter| {
                waiter.send(ImageReady(clone_arc(image)))
            }
            self.image_map.insert(url, @clone_arc(image));
            self.future_image_map.remove(url);
          }
          none => fail ~"storing an image that isn't in the future map"
        }
    }
}

#[test]
fn should_exit_on_request() {

    let mock_resource_task = do spawn_listener |from_client| {

        // infer me
        let from_client: port<resource_task::ControlMsg> = from_client;

        loop {
            match from_client.recv() {
              resource_task::Exit => break,
              _ => ()
            }
        }
    };

    let image_cache_task = image_cache_task(mock_resource_task);
    let url = make_url(~"file", none);

    image_cache_task.send(Exit);
    mock_resource_task.send(resource_task::Exit);
}

#[test]
#[should_fail]
fn should_fail_if_unprefetched_image_is_requested() {

    let mock_resource_task = do spawn_listener |from_client| {
    };

    let image_cache_task = image_cache_task(mock_resource_task);
    let url = make_url(~"file", none);

    let request = port();
    image_cache_task.send(GetImage(url, request.chan()));
    request.recv();
}

#[test]
fn should_request_url_from_resource_task_on_prefetch() {
    let url_requested = port();
    let url_requested_chan = url_requested.chan();

    let mock_resource_task = do spawn_listener |from_client| {

        // infer me
        let from_client: port<resource_task::ControlMsg> = from_client;

        loop {
            match from_client.recv() {
              resource_task::Load(url, _) => url_requested_chan.send(()),
              resource_task::Exit => break
            }
        }
    };

    let image_cache_task = image_cache_task(mock_resource_task);
    let url = make_url(~"file", none);

    image_cache_task.send(Prefetch(url));
    url_requested.recv();
    image_cache_task.send(Exit);
    mock_resource_task.send(resource_task::Exit);
}

#[test]
fn should_not_request_url_from_resource_task_on_multiple_prefetches() {
    let url_requested = port();
    let url_requested_chan = url_requested.chan();

    let mock_resource_task = do spawn_listener |from_client| {

        // infer me
        let from_client: port<resource_task::ControlMsg> = from_client;

        loop {
            match from_client.recv() {
              resource_task::Load(url, _) => url_requested_chan.send(()),
              resource_task::Exit => break
            }
        }
    };

    let image_cache_task = image_cache_task(mock_resource_task);
    let url = make_url(~"file", none);

    image_cache_task.send(Prefetch(url));
    image_cache_task.send(Prefetch(url));
    url_requested.recv();
    image_cache_task.send(Exit);
    mock_resource_task.send(resource_task::Exit);
    assert !url_requested.peek()
}

#[test]
fn should_return_image_not_ready_if_data_has_not_arrived() {
    let mock_resource_task = do spawn_listener |from_client| {

        // infer me
        let from_client: port<resource_task::ControlMsg> = from_client;

        loop {
            match from_client.recv() {
              resource_task::Exit => break,
              _ => ()
            }
        }
    };

    let image_cache_task = image_cache_task(mock_resource_task);
    let url = make_url(~"file", none);

    image_cache_task.send(Prefetch(url));
    let response_port = port();
    image_cache_task.send(GetImage(url, response_port.chan()));
    assert response_port.recv() == ImageNotReady;
    image_cache_task.send(Exit);
    mock_resource_task.send(resource_task::Exit);
}

#[test]
fn should_return_decoded_image_data_if_data_has_arrived() {

    let image_bin_sent = port();
    let image_bin_sent_chan = image_bin_sent.chan();

    let mock_resource_task = do spawn_listener |from_client| {

        // infer me
        let from_client: port<resource_task::ControlMsg> = from_client;

        loop {
            match from_client.recv() {
              resource_task::Load(_, response) => {
                response.send(resource_task::Payload(test_image_bin()));
                response.send(resource_task::Done(result::ok(())));
                image_bin_sent_chan.send(());
              }
              resource_task::Exit => break
            }
        }
    };

    let image_cache_task = image_cache_task(mock_resource_task);
    let url = make_url(~"file", none);

    image_cache_task.send(Prefetch(url));

    // Wait until our mock resource task has sent the image to the image cache
    image_bin_sent.recv();

    let response_port = port();
    image_cache_task.send(GetImage(url, response_port.chan()));
    match response_port.recv() {
      ImageReady(_) => (),
      _ => fail
    }

    image_cache_task.send(Exit);
    mock_resource_task.send(resource_task::Exit);
}

#[test]
fn should_return_decoded_image_data_for_multiple_requests() {

    let image_bin_sent = port();
    let image_bin_sent_chan = image_bin_sent.chan();

    let mock_resource_task = do spawn_listener |from_client| {

        // infer me
        let from_client: port<resource_task::ControlMsg> = from_client;

        loop {
            match from_client.recv() {
              resource_task::Load(_, response) => {
                response.send(resource_task::Payload(test_image_bin()));
                response.send(resource_task::Done(result::ok(())));
                image_bin_sent_chan.send(());
              }
              resource_task::Exit => break
            }
        }
    };

    let image_cache_task = image_cache_task(mock_resource_task);
    let url = make_url(~"file", none);

    image_cache_task.send(Prefetch(url));

    // Wait until our mock resource task has sent the image to the image cache
    image_bin_sent.recv();

    for iter::repeat(2) {
        let response_port = port();
        image_cache_task.send(GetImage(url, response_port.chan()));
        match response_port.recv() {
          ImageReady(_) => (),
          _ => fail
        }
    }

    image_cache_task.send(Exit);
    mock_resource_task.send(resource_task::Exit);
}

#[test]
fn should_not_request_image_from_resource_task_if_image_is_already_available() {

    let image_bin_sent = port();
    let image_bin_sent_chan = image_bin_sent.chan();

    let resource_task_exited = port();
    let resource_task_exited_chan = resource_task_exited.chan();

    let mock_resource_task = do spawn_listener |from_client| {

        // infer me
        let from_client: port<resource_task::ControlMsg> = from_client;

        loop {
            match from_client.recv() {
              resource_task::Load(_, response) => {
                response.send(resource_task::Payload(test_image_bin()));
                response.send(resource_task::Done(result::ok(())));
                image_bin_sent_chan.send(());
              }
              resource_task::Exit => {
                resource_task_exited_chan.send(());
                break
              }
            }
        }
    };

    let image_cache_task = image_cache_task(mock_resource_task);
    let url = make_url(~"file", none);

    image_cache_task.send(Prefetch(url));

    // Wait until our mock resource task has sent the image to the image cache
    image_bin_sent.recv();

    let response_port = port();
    image_cache_task.send(GetImage(url, response_port.chan()));
    match response_port.recv() {
      ImageReady(_) => (),
      _ => fail
    }

    image_cache_task.send(Prefetch(url));

    let response_port = port();
    image_cache_task.send(GetImage(url, response_port.chan()));
    response_port.recv();

    image_cache_task.send(Exit);
    mock_resource_task.send(resource_task::Exit);

    resource_task_exited.recv();

    // Our resource task should not have received another request for the image
    // because it's already cached
    assert !image_bin_sent.peek();
}

#[test]
fn should_return_not_ready_if_image_bin_cannot_be_fetched() {

    let image_bin_sent = port();
    let image_bin_sent_chan = image_bin_sent.chan();

    let mock_resource_task = do spawn_listener |from_client| {

        // infer me
        let from_client: port<resource_task::ControlMsg> = from_client;

        loop {
            match from_client.recv() {
              resource_task::Load(_, response) => {
                response.send(resource_task::Payload(test_image_bin()));
                // ERROR fetching image
                response.send(resource_task::Done(result::err(())));
                image_bin_sent_chan.send(());
              }
              resource_task::Exit => break
            }
        }
    };

    let image_cache_task = image_cache_task(mock_resource_task);
    let url = make_url(~"file", none);

    image_cache_task.send(Prefetch(url));

    // Wait until our mock resource task has sent the image to the image cache
    image_bin_sent.recv();

    let response_port = port();
    image_cache_task.send(GetImage(url, response_port.chan()));
    match response_port.recv() {
      ImageNotReady => (),
      _ => fail
    }

    image_cache_task.send(Exit);
    mock_resource_task.send(resource_task::Exit);
}
<commit_msg>Add docs to image_cache_task<commit_after>export Msg, Prefetch, GetImage, Exit;
export ImageResponseMsg, ImageReady, ImageNotReady;
export ImageCacheTask;
export image_cache_task;

import image::base::{Image, load_from_memory, test_image_bin};
import std::net::url::url;
import util::url::{make_url, UrlMap, url_map};
import comm::{chan, port};
import task::{spawn, spawn_listener};
import resource::resource_task;
import resource_task::ResourceTask;
import core::arc::arc;
import clone_arc = core::arc::clone;

enum Msg {
    /// Tell the cache that we may need a particular image soon. Must be posted
    /// before GetImage
    Prefetch(url),
    /// Request an Image object for a URL
    GetImage(url, chan<ImageResponseMsg>),
    /// Used by the decoder tasks to post decoded images back to the cache
    StoreImage(url, arc<~Image>),
    Exit
}

enum ImageResponseMsg {
    ImageReady(arc<~Image>),
    ImageNotReady
}

type ImageCacheTask = chan<Msg>;

fn image_cache_task(resource_task: ResourceTask) -> ImageCacheTask {
    do spawn_listener |from_client| {
        ImageCache {
            resource_task: resource_task,
            from_client: from_client,
            prefetch_map: url_map(),
            future_image_map: url_map(),
            image_map: url_map()
        }.run();
    }
}

struct ImageCache {
    /// A handle to the resource task for fetching the image binaries
    resource_task: ResourceTask;
    /// The port on which we'll receive client requests
    from_client: port<Msg>;
    /// A map from URLs to the partially loaded compressed image data.
    /// Once the data is complete it is then sent to a decoder
    prefetch_map: UrlMap<@PrefetchData>;
    /// A list of clients waiting on images that are currently being decoded
    future_image_map: UrlMap<@FutureData>;
    /// The cache of decoded images
    image_map: UrlMap<@arc<~Image>>;
}

struct PrefetchData {
    response_port: port<resource_task::ProgressMsg>;
    mut data: ~[u8];
}

struct FutureData {
    mut waiters: ~[chan<ImageResponseMsg>];
}

impl ImageCache {

    fn run() {

        loop {
            match self.from_client.recv() {
              Prefetch(url) => self.prefetch(url),
              GetImage(url, response) => self.get_image(url, response),
              StoreImage(url, image) => self.store_image(url, &image),
              Exit => break
            }
        }
    }

    /*priv*/ fn prefetch(url: url) {
        if self.image_map.contains_key(url) {
            // We've already decoded this image
            return
        }

        if self.future_image_map.contains_key(url) {
            // We've already begun decoding this image
            return
        }

        if self.prefetch_map.contains_key(url) {
            // We're already waiting for this image
            return
        }

        let response_port = port();
        self.resource_task.send(resource_task::Load(url, response_port.chan()));

        let prefetch_data = @PrefetchData {
            response_port: response_port,
            data: ~[]
        };

        self.prefetch_map.insert(url, prefetch_data);
    }

    /*priv*/ fn get_image(url: url, response: chan<ImageResponseMsg>) {
        match self.image_map.find(url) {
          some(image) => {
            response.send(ImageReady(clone_arc(image)));
            return
          }
          none => ()
        }

        match self.future_image_map.find(url) {
          some(future_data) => {
            // We've started decoding this image but haven't recieved it back yet.
            // Put this client on the wait list
            vec::push(future_data.waiters, response);
            return
          }
          none => ()
        }

        match self.prefetch_map.find(url) {
          some(prefetch_data) => {

            let mut image_sent = false;

            while prefetch_data.response_port.peek() {
                match prefetch_data.response_port.recv() {
                  resource_task::Payload(data) => {
                    prefetch_data.data += data;
                  }
                  resource_task::Done(result::ok(*)) => {
                    // We've got the entire image binary
                    let mut data = ~[];
                    data <-> prefetch_data.data;
                    let data <- data; // freeze for capture

                    let to_cache = self.from_client.chan();

                    do spawn {
                        let image = arc(~load_from_memory(data));
                        // Send the image to the original requester
                        response.send(ImageReady(clone_arc(&image)));
                        to_cache.send(StoreImage(url, clone_arc(&image)));
                    }

                    let future_data = @FutureData {
                        waiters: ~[]
                    };

                    self.prefetch_map.remove(url);
                    self.future_image_map.insert(url, future_data);

                    image_sent = true;
                    break;
                  }
                  resource_task::Done(result::err(*)) => {
                    // FIXME: need to actually report the failed image load
                    self.prefetch_map.remove(url);
                    break;
                  }
                }
            }

            if !image_sent {
                response.send(ImageNotReady);
            }
          }
          none => fail ~"got a request for image data without prefetch"
        }
    }

    /*priv*/ fn store_image(url: url, image: &arc<~Image>) {
        match self.future_image_map.find(url) {
          some(future_data) => {

            let mut waiters = ~[];
            waiters <-> future_data.waiters;

            // Send the image to all those who requested it while
            // it was being decoded
            for waiters.each |waiter| {
                waiter.send(ImageReady(clone_arc(image)))
            }
            self.image_map.insert(url, @clone_arc(image));
            self.future_image_map.remove(url);
          }
          none => fail ~"storing an image that isn't in the future map"
        }
    }
}

#[test]
fn should_exit_on_request() {

    let mock_resource_task = do spawn_listener |from_client| {

        // infer me
        let from_client: port<resource_task::ControlMsg> = from_client;

        loop {
            match from_client.recv() {
              resource_task::Exit => break,
              _ => ()
            }
        }
    };

    let image_cache_task = image_cache_task(mock_resource_task);
    let url = make_url(~"file", none);

    image_cache_task.send(Exit);
    mock_resource_task.send(resource_task::Exit);
}

#[test]
#[should_fail]
fn should_fail_if_unprefetched_image_is_requested() {

    let mock_resource_task = do spawn_listener |from_client| {
    };

    let image_cache_task = image_cache_task(mock_resource_task);
    let url = make_url(~"file", none);

    let request = port();
    image_cache_task.send(GetImage(url, request.chan()));
    request.recv();
}

#[test]
fn should_request_url_from_resource_task_on_prefetch() {
    let url_requested = port();
    let url_requested_chan = url_requested.chan();

    let mock_resource_task = do spawn_listener |from_client| {

        // infer me
        let from_client: port<resource_task::ControlMsg> = from_client;

        loop {
            match from_client.recv() {
              resource_task::Load(url, _) => url_requested_chan.send(()),
              resource_task::Exit => break
            }
        }
    };

    let image_cache_task = image_cache_task(mock_resource_task);
    let url = make_url(~"file", none);

    image_cache_task.send(Prefetch(url));
    url_requested.recv();
    image_cache_task.send(Exit);
    mock_resource_task.send(resource_task::Exit);
}

#[test]
fn should_not_request_url_from_resource_task_on_multiple_prefetches() {
    let url_requested = port();
    let url_requested_chan = url_requested.chan();

    let mock_resource_task = do spawn_listener |from_client| {

        // infer me
        let from_client: port<resource_task::ControlMsg> = from_client;

        loop {
            match from_client.recv() {
              resource_task::Load(url, _) => url_requested_chan.send(()),
              resource_task::Exit => break
            }
        }
    };

    let image_cache_task = image_cache_task(mock_resource_task);
    let url = make_url(~"file", none);

    image_cache_task.send(Prefetch(url));
    image_cache_task.send(Prefetch(url));
    url_requested.recv();
    image_cache_task.send(Exit);
    mock_resource_task.send(resource_task::Exit);
    assert !url_requested.peek()
}

#[test]
fn should_return_image_not_ready_if_data_has_not_arrived() {
    let mock_resource_task = do spawn_listener |from_client| {

        // infer me
        let from_client: port<resource_task::ControlMsg> = from_client;

        loop {
            match from_client.recv() {
              resource_task::Exit => break,
              _ => ()
            }
        }
    };

    let image_cache_task = image_cache_task(mock_resource_task);
    let url = make_url(~"file", none);

    image_cache_task.send(Prefetch(url));
    let response_port = port();
    image_cache_task.send(GetImage(url, response_port.chan()));
    assert response_port.recv() == ImageNotReady;
    image_cache_task.send(Exit);
    mock_resource_task.send(resource_task::Exit);
}

#[test]
fn should_return_decoded_image_data_if_data_has_arrived() {

    let image_bin_sent = port();
    let image_bin_sent_chan = image_bin_sent.chan();

    let mock_resource_task = do spawn_listener |from_client| {

        // infer me
        let from_client: port<resource_task::ControlMsg> = from_client;

        loop {
            match from_client.recv() {
              resource_task::Load(_, response) => {
                response.send(resource_task::Payload(test_image_bin()));
                response.send(resource_task::Done(result::ok(())));
                image_bin_sent_chan.send(());
              }
              resource_task::Exit => break
            }
        }
    };

    let image_cache_task = image_cache_task(mock_resource_task);
    let url = make_url(~"file", none);

    image_cache_task.send(Prefetch(url));

    // Wait until our mock resource task has sent the image to the image cache
    image_bin_sent.recv();

    let response_port = port();
    image_cache_task.send(GetImage(url, response_port.chan()));
    match response_port.recv() {
      ImageReady(_) => (),
      _ => fail
    }

    image_cache_task.send(Exit);
    mock_resource_task.send(resource_task::Exit);
}

#[test]
fn should_return_decoded_image_data_for_multiple_requests() {

    let image_bin_sent = port();
    let image_bin_sent_chan = image_bin_sent.chan();

    let mock_resource_task = do spawn_listener |from_client| {

        // infer me
        let from_client: port<resource_task::ControlMsg> = from_client;

        loop {
            match from_client.recv() {
              resource_task::Load(_, response) => {
                response.send(resource_task::Payload(test_image_bin()));
                response.send(resource_task::Done(result::ok(())));
                image_bin_sent_chan.send(());
              }
              resource_task::Exit => break
            }
        }
    };

    let image_cache_task = image_cache_task(mock_resource_task);
    let url = make_url(~"file", none);

    image_cache_task.send(Prefetch(url));

    // Wait until our mock resource task has sent the image to the image cache
    image_bin_sent.recv();

    for iter::repeat(2) {
        let response_port = port();
        image_cache_task.send(GetImage(url, response_port.chan()));
        match response_port.recv() {
          ImageReady(_) => (),
          _ => fail
        }
    }

    image_cache_task.send(Exit);
    mock_resource_task.send(resource_task::Exit);
}

#[test]
fn should_not_request_image_from_resource_task_if_image_is_already_available() {

    let image_bin_sent = port();
    let image_bin_sent_chan = image_bin_sent.chan();

    let resource_task_exited = port();
    let resource_task_exited_chan = resource_task_exited.chan();

    let mock_resource_task = do spawn_listener |from_client| {

        // infer me
        let from_client: port<resource_task::ControlMsg> = from_client;

        loop {
            match from_client.recv() {
              resource_task::Load(_, response) => {
                response.send(resource_task::Payload(test_image_bin()));
                response.send(resource_task::Done(result::ok(())));
                image_bin_sent_chan.send(());
              }
              resource_task::Exit => {
                resource_task_exited_chan.send(());
                break
              }
            }
        }
    };

    let image_cache_task = image_cache_task(mock_resource_task);
    let url = make_url(~"file", none);

    image_cache_task.send(Prefetch(url));

    // Wait until our mock resource task has sent the image to the image cache
    image_bin_sent.recv();

    let response_port = port();
    image_cache_task.send(GetImage(url, response_port.chan()));
    match response_port.recv() {
      ImageReady(_) => (),
      _ => fail
    }

    image_cache_task.send(Prefetch(url));

    let response_port = port();
    image_cache_task.send(GetImage(url, response_port.chan()));
    response_port.recv();

    image_cache_task.send(Exit);
    mock_resource_task.send(resource_task::Exit);

    resource_task_exited.recv();

    // Our resource task should not have received another request for the image
    // because it's already cached
    assert !image_bin_sent.peek();
}

#[test]
fn should_return_not_ready_if_image_bin_cannot_be_fetched() {

    let image_bin_sent = port();
    let image_bin_sent_chan = image_bin_sent.chan();

    let mock_resource_task = do spawn_listener |from_client| {

        // infer me
        let from_client: port<resource_task::ControlMsg> = from_client;

        loop {
            match from_client.recv() {
              resource_task::Load(_, response) => {
                response.send(resource_task::Payload(test_image_bin()));
                // ERROR fetching image
                response.send(resource_task::Done(result::err(())));
                image_bin_sent_chan.send(());
              }
              resource_task::Exit => break
            }
        }
    };

    let image_cache_task = image_cache_task(mock_resource_task);
    let url = make_url(~"file", none);

    image_cache_task.send(Prefetch(url));

    // Wait until our mock resource task has sent the image to the image cache
    image_bin_sent.recv();

    let response_port = port();
    image_cache_task.send(GetImage(url, response_port.chan()));
    match response_port.recv() {
      ImageNotReady => (),
      _ => fail
    }

    image_cache_task.send(Exit);
    mock_resource_task.send(resource_task::Exit);
}
<|endoftext|>