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

use font_template::{FontTemplate, FontTemplateDescriptor};
use ipc_channel::ipc::{self, IpcReceiver, IpcSender};
use ipc_channel::router::ROUTER;
use mime::{TopLevel, SubLevel};
use net_traits::{AsyncResponseTarget, LoadContext, PendingAsyncLoad, ResourceTask, ResponseAction};
use platform::font_context::FontContextHandle;
use platform::font_list::for_each_available_family;
use platform::font_list::for_each_variation;
use platform::font_list::last_resort_font_families;
use platform::font_list::system_default_family;
use platform::font_template::FontTemplateData;
use std::borrow::ToOwned;
use std::collections::HashMap;
use std::mem;
use std::sync::mpsc::channel;
use std::sync::{Arc, Mutex};
use string_cache::Atom;
use style::font_face::Source;
use style::properties::longhands::font_family::computed_value::FontFamily;
use url::Url;
use util::str::LowercaseString;
use util::task::spawn_named;

/// A list of font templates that make up a given font family.
struct FontTemplates {
    templates: Vec<FontTemplate>,
}

impl FontTemplates {
    fn new() -> FontTemplates {
        FontTemplates {
            templates: vec!(),
        }
    }

    /// Find a font in this family that matches a given descriptor.
    fn find_font_for_style<'a>(&'a mut self, desc: &FontTemplateDescriptor, fctx: &FontContextHandle)
                               -> Option<Arc<FontTemplateData>> {
        // TODO(Issue #189): optimize lookup for
        // regular/bold/italic/bolditalic with fixed offsets and a
        // static decision table for fallback between these values.

        // TODO(Issue #190): if not in the fast path above, do
        // expensive matching of weights, etc.
        for template in &mut self.templates {
            let maybe_template = template.data_for_descriptor(fctx, desc);
            if maybe_template.is_some() {
                return maybe_template;
            }
        }

        // If a request is made for a font family that exists,
        // pick the first valid font in the family if we failed
        // to find an exact match for the descriptor.
        for template in &mut self.templates {
            let maybe_template = template.get();
            if maybe_template.is_some() {
                return maybe_template;
            }
        }

        None
    }

    fn add_template(&mut self, identifier: Atom, maybe_data: Option<Vec<u8>>) {
        for template in &self.templates {
            if *template.identifier() == identifier {
                return;
            }
        }

        let template = FontTemplate::new(identifier, maybe_data);
        self.templates.push(template);
    }
}

/// Commands that the FontContext sends to the font cache task.
#[derive(Deserialize, Serialize)]
pub enum Command {
    GetFontTemplate(FontFamily, FontTemplateDescriptor, IpcSender<Reply>),
    GetLastResortFontTemplate(FontTemplateDescriptor, IpcSender<Reply>),
    AddWebFont(FontFamily, Source, IpcSender<()>),
    AddDownloadedWebFont(FontFamily, Url, Vec<u8>, IpcSender<()>),
    Exit(IpcSender<()>),
}

/// Reply messages sent from the font cache task to the FontContext caller.
#[derive(Deserialize, Serialize)]
pub enum Reply {
    GetFontTemplateReply(Option<Arc<FontTemplateData>>),
}

/// The font cache task itself. It maintains a list of reference counted
/// font templates that are currently in use.
struct FontCache {
    port: IpcReceiver<Command>,
    channel_to_self: IpcSender<Command>,
    generic_fonts: HashMap<FontFamily, LowercaseString>,
    local_families: HashMap<LowercaseString, FontTemplates>,
    web_families: HashMap<LowercaseString, FontTemplates>,
    font_context: FontContextHandle,
    resource_task: ResourceTask,
}

fn populate_generic_fonts() -> HashMap<FontFamily, LowercaseString> {
    let mut generic_fonts = HashMap::with_capacity(5);

    append_map(&mut generic_fonts, FontFamily::Serif, "Times New Roman");
    append_map(&mut generic_fonts, FontFamily::SansSerif, "Arial");
    append_map(&mut generic_fonts, FontFamily::Cursive, "Apple Chancery");
    append_map(&mut generic_fonts, FontFamily::Fantasy, "Papyrus");
    append_map(&mut generic_fonts, FontFamily::Monospace, "Menlo");

    fn append_map(generic_fonts: &mut HashMap<FontFamily, LowercaseString>,
                  font_family: FontFamily,
                  mapped_name: &str) {
        let family_name = {
            let opt_system_default = system_default_family(font_family.name());
            match opt_system_default {
                Some(system_default) => LowercaseString::new(&system_default),
                None => LowercaseString::new(mapped_name)
            }
        };

        generic_fonts.insert(font_family, family_name);
    }

    generic_fonts
}

impl FontCache {
    fn run(&mut self) {
        loop {
            let msg = self.port.recv().unwrap();

            match msg {
                Command::GetFontTemplate(family, descriptor, result) => {
                    let maybe_font_template = self.find_font_template(&family, &descriptor);
                    result.send(Reply::GetFontTemplateReply(maybe_font_template)).unwrap();
                }
                Command::GetLastResortFontTemplate(descriptor, result) => {
                    let font_template = self.last_resort_font_template(&descriptor);
                    result.send(Reply::GetFontTemplateReply(Some(font_template))).unwrap();
                }
                Command::AddWebFont(family, src, result) => {
                    let family_name = LowercaseString::new(family.name());
                    if !self.web_families.contains_key(&family_name) {
                        let templates = FontTemplates::new();
                        self.web_families.insert(family_name, templates);
                    }

                    match src {
                        Source::Url(ref url_source) => {
                            let url = &url_source.url;
                            let load = PendingAsyncLoad::new(LoadContext::Font,
                                                             self.resource_task.clone(),
                                                             url.clone(),
                                                             None);
                            let (data_sender, data_receiver) = ipc::channel().unwrap();
                            let data_target = AsyncResponseTarget {
                                sender: data_sender,
                            };
                            load.load_async(data_target);
                            let channel_to_self = self.channel_to_self.clone();
                            let url = (*url).clone();
                            let bytes = Mutex::new(Vec::new());
                            let response_valid = Mutex::new(false);
                            ROUTER.add_route(data_receiver.to_opaque(), box move |message| {
                                let response: ResponseAction = message.to().unwrap();
                                match response {
                                    ResponseAction::HeadersAvailable(metadata) => {
                                        let is_response_valid =
                                            metadata.content_type.as_ref().map_or(false, |content_type| {
                                                let mime = &content_type.0;
                                                is_supported_font_type(&mime.0, &mime.1)
                                            });
                                        info!("{} font with MIME type {:?}",
                                              if is_response_valid { "Loading" } else { "Ignoring" },
                                              metadata.content_type);
                                        *response_valid.lock().unwrap() = is_response_valid;
                                    }
                                    ResponseAction::ResponseComplete(Err(_)) => {}
                                    ResponseAction::DataAvailable(new_bytes) => {
                                        if *response_valid.lock().unwrap() {
                                            bytes.lock().unwrap().extend(new_bytes.into_iter())
                                        }
                                    }
                                    ResponseAction::ResponseComplete(Ok(_)) => {
                                        if !*response_valid.lock().unwrap() {
                                            return;
                                        }
                                        let mut bytes = bytes.lock().unwrap();
                                        let bytes = mem::replace(&mut *bytes, Vec::new());
                                        let command =
                                            Command::AddDownloadedWebFont(family.clone(),
                                                                          url.clone(),
                                                                          bytes,
                                                                          result.clone());
                                        channel_to_self.send(command).unwrap();
                                    }
                                }
                            });
                        }
                        Source::Local(ref family) => {
                            let family_name = LowercaseString::new(family.name());
                            let templates = &mut self.web_families.get_mut(&family_name).unwrap();
                            for_each_variation(&family_name, |path| {
                                templates.add_template(Atom::from(&*path), None);
                            });
                            result.send(()).unwrap();
                        }
                    }
                }
                Command::AddDownloadedWebFont(family, url, bytes, result) => {
                    let family_name = LowercaseString::new(family.name());

                    let templates = &mut self.web_families.get_mut(&family_name).unwrap();
                    templates.add_template(Atom::from(&*url.to_string()), Some(bytes));
                    drop(result.send(()));
                }
                Command::Exit(result) => {
                    result.send(()).unwrap();
                    break;
                }
            }
        }
    }

    fn refresh_local_families(&mut self) {
        self.local_families.clear();
        for_each_available_family(|family_name| {
            let family_name = LowercaseString::new(&family_name);
            if !self.local_families.contains_key(&family_name) {
                let templates = FontTemplates::new();
                self.local_families.insert(family_name, templates);
            }
        });
    }

    fn transform_family(&self, family: &FontFamily) -> LowercaseString {
        match self.generic_fonts.get(family) {
            None => LowercaseString::new(family.name()),
            Some(mapped_family) => (*mapped_family).clone()
        }
    }

    fn find_font_in_local_family<'a>(&'a mut self, family_name: &LowercaseString, desc: &FontTemplateDescriptor)
                                -> Option<Arc<FontTemplateData>> {
        // TODO(Issue #188): look up localized font family names if canonical name not found
        // look up canonical name
        if self.local_families.contains_key(family_name) {
            debug!("FontList: Found font family with name={}", &**family_name);
            let s = self.local_families.get_mut(family_name).unwrap();

            if s.templates.is_empty() {
                for_each_variation(family_name, |path| {
                    s.add_template(Atom::from(&*path), None);
                });
            }

            // TODO(Issue #192: handle generic font families, like 'serif' and 'sans-serif'.
            // if such family exists, try to match style to a font

            s.find_font_for_style(desc, &self.font_context)
        } else {
            debug!("FontList: Couldn't find font family with name={}", &**family_name);
            None
        }
    }

    fn find_font_in_web_family<'a>(&'a mut self, family: &FontFamily, desc: &FontTemplateDescriptor)
                                -> Option<Arc<FontTemplateData>> {
        let family_name = LowercaseString::new(family.name());

        if self.web_families.contains_key(&family_name) {
            let templates = self.web_families.get_mut(&family_name).unwrap();
            let maybe_font = templates.find_font_for_style(desc, &self.font_context);
            maybe_font
        } else {
            None
        }
    }

    fn find_font_template(&mut self, family: &FontFamily, desc: &FontTemplateDescriptor)
                            -> Option<Arc<FontTemplateData>> {
        self.find_font_in_web_family(family, desc)
            .or_else(|| {
                let transformed_family = self.transform_family(family);
                self.find_font_in_local_family(&transformed_family, desc)
            })
    }

    fn last_resort_font_template(&mut self, desc: &FontTemplateDescriptor)
                                        -> Arc<FontTemplateData> {
        let last_resort = last_resort_font_families();

        for family in &last_resort {
            let family = LowercaseString::new(family);
            let maybe_font_in_family = self.find_font_in_local_family(&family, desc);
            if let Some(family) = maybe_font_in_family {
                return family;
            }
        }

        panic!("Unable to find any fonts that match (do you have fallback fonts installed?)");
    }
}

/// The public interface to the font cache task, used exclusively by
/// the per-thread/task FontContext structures.
#[derive(Clone, Deserialize, Serialize)]
pub struct FontCacheTask {
    chan: IpcSender<Command>,
}

impl FontCacheTask {
    pub fn new(resource_task: ResourceTask) -> FontCacheTask {
        let (chan, port) = ipc::channel().unwrap();

        let channel_to_self = chan.clone();
        spawn_named("FontCacheTask".to_owned(), move || {
            // TODO: Allow users to specify these.
            let generic_fonts = populate_generic_fonts();

            let mut cache = FontCache {
                port: port,
                channel_to_self: channel_to_self,
                generic_fonts: generic_fonts,
                local_families: HashMap::new(),
                web_families: HashMap::new(),
                font_context: FontContextHandle::new(),
                resource_task: resource_task,
            };

            cache.refresh_local_families();
            cache.run();
        });

        FontCacheTask {
            chan: chan,
        }
    }

    pub fn find_font_template(&self, family: FontFamily, desc: FontTemplateDescriptor)
                                                -> Option<Arc<FontTemplateData>> {

        let (response_chan, response_port) = ipc::channel().unwrap();
        self.chan.send(Command::GetFontTemplate(family, desc, response_chan)).unwrap();

        let reply = response_port.recv().unwrap();

        match reply {
            Reply::GetFontTemplateReply(data) => {
                data
            }
        }
    }

    pub fn last_resort_font_template(&self, desc: FontTemplateDescriptor)
                                                -> Arc<FontTemplateData> {

        let (response_chan, response_port) = ipc::channel().unwrap();
        self.chan.send(Command::GetLastResortFontTemplate(desc, response_chan)).unwrap();

        let reply = response_port.recv().unwrap();

        match reply {
            Reply::GetFontTemplateReply(data) => {
                data.unwrap()
            }
        }
    }

    pub fn add_web_font(&self, family: FontFamily, src: Source, sender: IpcSender<()>) {
        self.chan.send(Command::AddWebFont(family, src, sender)).unwrap();
    }

    pub fn exit(&self) {
        let (response_chan, response_port) = ipc::channel().unwrap();
        self.chan.send(Command::Exit(response_chan)).unwrap();
        response_port.recv().unwrap();
    }
}

// derived from http://stackoverflow.com/a/10864297/3830
fn is_supported_font_type(toplevel: &TopLevel, sublevel: &SubLevel) -> bool {
    match (toplevel, sublevel) {
        (&TopLevel::Application, &SubLevel::Ext(ref ext)) => {
            match &ext[..] {
                //FIXME: once sniffing is enabled by default, we shouldn't need nonstandard
                //       MIME types here.
                "font-sfnt" | "x-font-ttf" | "x-font-truetype" | "x-font-opentype" => true,
                _ => false,
            }
        }
        _ => false,
    }
}
<commit_msg>Send a response to font load initiators even when a font is ignored.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use font_template::{FontTemplate, FontTemplateDescriptor};
use ipc_channel::ipc::{self, IpcReceiver, IpcSender};
use ipc_channel::router::ROUTER;
use mime::{TopLevel, SubLevel};
use net_traits::{AsyncResponseTarget, LoadContext, PendingAsyncLoad, ResourceTask, ResponseAction};
use platform::font_context::FontContextHandle;
use platform::font_list::for_each_available_family;
use platform::font_list::for_each_variation;
use platform::font_list::last_resort_font_families;
use platform::font_list::system_default_family;
use platform::font_template::FontTemplateData;
use std::borrow::ToOwned;
use std::collections::HashMap;
use std::mem;
use std::sync::mpsc::channel;
use std::sync::{Arc, Mutex};
use string_cache::Atom;
use style::font_face::Source;
use style::properties::longhands::font_family::computed_value::FontFamily;
use url::Url;
use util::str::LowercaseString;
use util::task::spawn_named;

/// A list of font templates that make up a given font family.
struct FontTemplates {
    templates: Vec<FontTemplate>,
}

impl FontTemplates {
    fn new() -> FontTemplates {
        FontTemplates {
            templates: vec!(),
        }
    }

    /// Find a font in this family that matches a given descriptor.
    fn find_font_for_style<'a>(&'a mut self, desc: &FontTemplateDescriptor, fctx: &FontContextHandle)
                               -> Option<Arc<FontTemplateData>> {
        // TODO(Issue #189): optimize lookup for
        // regular/bold/italic/bolditalic with fixed offsets and a
        // static decision table for fallback between these values.

        // TODO(Issue #190): if not in the fast path above, do
        // expensive matching of weights, etc.
        for template in &mut self.templates {
            let maybe_template = template.data_for_descriptor(fctx, desc);
            if maybe_template.is_some() {
                return maybe_template;
            }
        }

        // If a request is made for a font family that exists,
        // pick the first valid font in the family if we failed
        // to find an exact match for the descriptor.
        for template in &mut self.templates {
            let maybe_template = template.get();
            if maybe_template.is_some() {
                return maybe_template;
            }
        }

        None
    }

    fn add_template(&mut self, identifier: Atom, maybe_data: Option<Vec<u8>>) {
        for template in &self.templates {
            if *template.identifier() == identifier {
                return;
            }
        }

        let template = FontTemplate::new(identifier, maybe_data);
        self.templates.push(template);
    }
}

/// Commands that the FontContext sends to the font cache task.
#[derive(Deserialize, Serialize)]
pub enum Command {
    GetFontTemplate(FontFamily, FontTemplateDescriptor, IpcSender<Reply>),
    GetLastResortFontTemplate(FontTemplateDescriptor, IpcSender<Reply>),
    AddWebFont(FontFamily, Source, IpcSender<()>),
    AddDownloadedWebFont(FontFamily, Url, Vec<u8>, IpcSender<()>),
    Exit(IpcSender<()>),
}

/// Reply messages sent from the font cache task to the FontContext caller.
#[derive(Deserialize, Serialize)]
pub enum Reply {
    GetFontTemplateReply(Option<Arc<FontTemplateData>>),
}

/// The font cache task itself. It maintains a list of reference counted
/// font templates that are currently in use.
struct FontCache {
    port: IpcReceiver<Command>,
    channel_to_self: IpcSender<Command>,
    generic_fonts: HashMap<FontFamily, LowercaseString>,
    local_families: HashMap<LowercaseString, FontTemplates>,
    web_families: HashMap<LowercaseString, FontTemplates>,
    font_context: FontContextHandle,
    resource_task: ResourceTask,
}

fn populate_generic_fonts() -> HashMap<FontFamily, LowercaseString> {
    let mut generic_fonts = HashMap::with_capacity(5);

    append_map(&mut generic_fonts, FontFamily::Serif, "Times New Roman");
    append_map(&mut generic_fonts, FontFamily::SansSerif, "Arial");
    append_map(&mut generic_fonts, FontFamily::Cursive, "Apple Chancery");
    append_map(&mut generic_fonts, FontFamily::Fantasy, "Papyrus");
    append_map(&mut generic_fonts, FontFamily::Monospace, "Menlo");

    fn append_map(generic_fonts: &mut HashMap<FontFamily, LowercaseString>,
                  font_family: FontFamily,
                  mapped_name: &str) {
        let family_name = {
            let opt_system_default = system_default_family(font_family.name());
            match opt_system_default {
                Some(system_default) => LowercaseString::new(&system_default),
                None => LowercaseString::new(mapped_name)
            }
        };

        generic_fonts.insert(font_family, family_name);
    }

    generic_fonts
}

impl FontCache {
    fn run(&mut self) {
        loop {
            let msg = self.port.recv().unwrap();

            match msg {
                Command::GetFontTemplate(family, descriptor, result) => {
                    let maybe_font_template = self.find_font_template(&family, &descriptor);
                    result.send(Reply::GetFontTemplateReply(maybe_font_template)).unwrap();
                }
                Command::GetLastResortFontTemplate(descriptor, result) => {
                    let font_template = self.last_resort_font_template(&descriptor);
                    result.send(Reply::GetFontTemplateReply(Some(font_template))).unwrap();
                }
                Command::AddWebFont(family, src, result) => {
                    let family_name = LowercaseString::new(family.name());
                    if !self.web_families.contains_key(&family_name) {
                        let templates = FontTemplates::new();
                        self.web_families.insert(family_name, templates);
                    }

                    match src {
                        Source::Url(ref url_source) => {
                            let url = &url_source.url;
                            let load = PendingAsyncLoad::new(LoadContext::Font,
                                                             self.resource_task.clone(),
                                                             url.clone(),
                                                             None);
                            let (data_sender, data_receiver) = ipc::channel().unwrap();
                            let data_target = AsyncResponseTarget {
                                sender: data_sender,
                            };
                            load.load_async(data_target);
                            let channel_to_self = self.channel_to_self.clone();
                            let url = (*url).clone();
                            let bytes = Mutex::new(Vec::new());
                            let response_valid = Mutex::new(false);
                            ROUTER.add_route(data_receiver.to_opaque(), box move |message| {
                                let response: ResponseAction = message.to().unwrap();
                                match response {
                                    ResponseAction::HeadersAvailable(metadata) => {
                                        let is_response_valid =
                                            metadata.content_type.as_ref().map_or(false, |content_type| {
                                                let mime = &content_type.0;
                                                is_supported_font_type(&mime.0, &mime.1)
                                            });
                                        info!("{} font with MIME type {:?}",
                                              if is_response_valid { "Loading" } else { "Ignoring" },
                                              metadata.content_type);
                                        *response_valid.lock().unwrap() = is_response_valid;
                                    }
                                    ResponseAction::ResponseComplete(Err(_)) => {}
                                    ResponseAction::DataAvailable(new_bytes) => {
                                        if *response_valid.lock().unwrap() {
                                            bytes.lock().unwrap().extend(new_bytes.into_iter())
                                        }
                                    }
                                    ResponseAction::ResponseComplete(Ok(_)) => {
                                        if !*response_valid.lock().unwrap() {
                                            drop(result.send(()));
                                            return;
                                        }
                                        let mut bytes = bytes.lock().unwrap();
                                        let bytes = mem::replace(&mut *bytes, Vec::new());
                                        let command =
                                            Command::AddDownloadedWebFont(family.clone(),
                                                                          url.clone(),
                                                                          bytes,
                                                                          result.clone());
                                        channel_to_self.send(command).unwrap();
                                    }
                                }
                            });
                        }
                        Source::Local(ref family) => {
                            let family_name = LowercaseString::new(family.name());
                            let templates = &mut self.web_families.get_mut(&family_name).unwrap();
                            for_each_variation(&family_name, |path| {
                                templates.add_template(Atom::from(&*path), None);
                            });
                            result.send(()).unwrap();
                        }
                    }
                }
                Command::AddDownloadedWebFont(family, url, bytes, result) => {
                    let family_name = LowercaseString::new(family.name());

                    let templates = &mut self.web_families.get_mut(&family_name).unwrap();
                    templates.add_template(Atom::from(&*url.to_string()), Some(bytes));
                    drop(result.send(()));
                }
                Command::Exit(result) => {
                    result.send(()).unwrap();
                    break;
                }
            }
        }
    }

    fn refresh_local_families(&mut self) {
        self.local_families.clear();
        for_each_available_family(|family_name| {
            let family_name = LowercaseString::new(&family_name);
            if !self.local_families.contains_key(&family_name) {
                let templates = FontTemplates::new();
                self.local_families.insert(family_name, templates);
            }
        });
    }

    fn transform_family(&self, family: &FontFamily) -> LowercaseString {
        match self.generic_fonts.get(family) {
            None => LowercaseString::new(family.name()),
            Some(mapped_family) => (*mapped_family).clone()
        }
    }

    fn find_font_in_local_family<'a>(&'a mut self, family_name: &LowercaseString, desc: &FontTemplateDescriptor)
                                -> Option<Arc<FontTemplateData>> {
        // TODO(Issue #188): look up localized font family names if canonical name not found
        // look up canonical name
        if self.local_families.contains_key(family_name) {
            debug!("FontList: Found font family with name={}", &**family_name);
            let s = self.local_families.get_mut(family_name).unwrap();

            if s.templates.is_empty() {
                for_each_variation(family_name, |path| {
                    s.add_template(Atom::from(&*path), None);
                });
            }

            // TODO(Issue #192: handle generic font families, like 'serif' and 'sans-serif'.
            // if such family exists, try to match style to a font

            s.find_font_for_style(desc, &self.font_context)
        } else {
            debug!("FontList: Couldn't find font family with name={}", &**family_name);
            None
        }
    }

    fn find_font_in_web_family<'a>(&'a mut self, family: &FontFamily, desc: &FontTemplateDescriptor)
                                -> Option<Arc<FontTemplateData>> {
        let family_name = LowercaseString::new(family.name());

        if self.web_families.contains_key(&family_name) {
            let templates = self.web_families.get_mut(&family_name).unwrap();
            let maybe_font = templates.find_font_for_style(desc, &self.font_context);
            maybe_font
        } else {
            None
        }
    }

    fn find_font_template(&mut self, family: &FontFamily, desc: &FontTemplateDescriptor)
                            -> Option<Arc<FontTemplateData>> {
        self.find_font_in_web_family(family, desc)
            .or_else(|| {
                let transformed_family = self.transform_family(family);
                self.find_font_in_local_family(&transformed_family, desc)
            })
    }

    fn last_resort_font_template(&mut self, desc: &FontTemplateDescriptor)
                                        -> Arc<FontTemplateData> {
        let last_resort = last_resort_font_families();

        for family in &last_resort {
            let family = LowercaseString::new(family);
            let maybe_font_in_family = self.find_font_in_local_family(&family, desc);
            if let Some(family) = maybe_font_in_family {
                return family;
            }
        }

        panic!("Unable to find any fonts that match (do you have fallback fonts installed?)");
    }
}

/// The public interface to the font cache task, used exclusively by
/// the per-thread/task FontContext structures.
#[derive(Clone, Deserialize, Serialize)]
pub struct FontCacheTask {
    chan: IpcSender<Command>,
}

impl FontCacheTask {
    pub fn new(resource_task: ResourceTask) -> FontCacheTask {
        let (chan, port) = ipc::channel().unwrap();

        let channel_to_self = chan.clone();
        spawn_named("FontCacheTask".to_owned(), move || {
            // TODO: Allow users to specify these.
            let generic_fonts = populate_generic_fonts();

            let mut cache = FontCache {
                port: port,
                channel_to_self: channel_to_self,
                generic_fonts: generic_fonts,
                local_families: HashMap::new(),
                web_families: HashMap::new(),
                font_context: FontContextHandle::new(),
                resource_task: resource_task,
            };

            cache.refresh_local_families();
            cache.run();
        });

        FontCacheTask {
            chan: chan,
        }
    }

    pub fn find_font_template(&self, family: FontFamily, desc: FontTemplateDescriptor)
                                                -> Option<Arc<FontTemplateData>> {

        let (response_chan, response_port) = ipc::channel().unwrap();
        self.chan.send(Command::GetFontTemplate(family, desc, response_chan)).unwrap();

        let reply = response_port.recv().unwrap();

        match reply {
            Reply::GetFontTemplateReply(data) => {
                data
            }
        }
    }

    pub fn last_resort_font_template(&self, desc: FontTemplateDescriptor)
                                                -> Arc<FontTemplateData> {

        let (response_chan, response_port) = ipc::channel().unwrap();
        self.chan.send(Command::GetLastResortFontTemplate(desc, response_chan)).unwrap();

        let reply = response_port.recv().unwrap();

        match reply {
            Reply::GetFontTemplateReply(data) => {
                data.unwrap()
            }
        }
    }

    pub fn add_web_font(&self, family: FontFamily, src: Source, sender: IpcSender<()>) {
        self.chan.send(Command::AddWebFont(family, src, sender)).unwrap();
    }

    pub fn exit(&self) {
        let (response_chan, response_port) = ipc::channel().unwrap();
        self.chan.send(Command::Exit(response_chan)).unwrap();
        response_port.recv().unwrap();
    }
}

// derived from http://stackoverflow.com/a/10864297/3830
fn is_supported_font_type(toplevel: &TopLevel, sublevel: &SubLevel) -> bool {
    match (toplevel, sublevel) {
        (&TopLevel::Application, &SubLevel::Ext(ref ext)) => {
            match &ext[..] {
                //FIXME: once sniffing is enabled by default, we shouldn't need nonstandard
                //       MIME types here.
                "font-sfnt" | "x-font-ttf" | "x-font-truetype" | "x-font-opentype" => true,
                _ => false,
            }
        }
        _ => false,
    }
}
<|endoftext|>


<commit_before>use internal::prelude::*;
use parking_lot::Mutex;
use std::collections::HashMap;
use std::sync::mpsc::{self, Sender};
use std::sync::Arc;
use std::thread;
use super::super::super::EventHandler;
use super::{
    ShardClientMessage,
    ShardId,
    ShardManagerMessage,
    ShardManagerMonitor,
    ShardQueuer,
    ShardQueuerMessage,
    ShardRunnerInfo,
};
use threadpool::ThreadPool;
use typemap::ShareMap;

#[cfg(feature = "framework")]
use framework::Framework;

/// A manager for handling the status of shards by starting them, restarting
/// them, and stopping them when required.
///
/// **Note**: The [`Client`] internally uses a shard manager. If you are using a
/// Client, then you do not need to make one of these.
///
/// # Examples
///
/// Initialize a shard manager with a framework responsible for shards 0 through
/// 2, of 5 total shards:
///
/// ```rust,no_run
/// extern crate parking_lot;
/// extern crate serenity;
/// extern crate threadpool;
/// extern crate typemap;
///
/// # use std::error::Error;
/// #
/// # #[cfg(feature = "framework")]
/// # fn try_main() -> Result<(), Box<Error>> {
/// #
/// use parking_lot::Mutex;
/// use serenity::client::bridge::gateway::ShardManager;
/// use serenity::client::EventHandler;
/// use serenity::http;
/// use std::sync::Arc;
/// use std::env;
/// use threadpool::ThreadPool;
/// use typemap::ShareMap;
///
/// struct Handler;
///
/// impl EventHandler for Handler { }
///
/// let token = env::var("DISCORD_TOKEN")?;
/// http::set_token(&token);
/// let token = Arc::new(Mutex::new(token));
///
/// let gateway_url = Arc::new(Mutex::new(http::get_gateway()?.url));
/// let data = Arc::new(Mutex::new(ShareMap::custom()));
/// let event_handler = Arc::new(Handler);
/// let framework = Arc::new(Mutex::new(None));
/// let threadpool = ThreadPool::with_name("my threadpool".to_owned(), 5);
///
/// ShardManager::new(
///     0, // the shard index to start initiating from
///     3, // the number of shards to initiate (this initiates 0, 1, and 2)
///     5, // the total number of shards in use
///     gateway_url,
///     token,
///     data,
///     event_handler,
///     framework,
///     threadpool,
/// );
/// #     Ok(())
/// # }
/// #
/// # #[cfg(not(feature = "framework"))]
/// # fn try_main() -> Result<(), Box<Error>> {
/// #     Ok(())
/// # }
/// #
/// # fn main() {
/// #     try_main().unwrap();
/// # }
/// ```
///
/// [`Client`]: ../../struct.Client.html
#[derive(Debug)]
pub struct ShardManager {
    /// The shard runners currently managed.
    runners: Arc<Mutex<HashMap<ShardId, ShardRunnerInfo>>>,
    /// The index of the first shard to initialize, 0-indexed.
    shard_index: u64,
    /// The number of shards to initialize.
    shard_init: u64,
    /// The total shards in use, 1-indexed.
    shard_total: u64,
    shard_queuer: Sender<ShardQueuerMessage>,
}

impl ShardManager {
    /// Creates a new shard manager, returning both the manager and a monitor
    /// for usage in a separate thread.
    #[cfg(feature = "framework")]
    #[cfg_attr(feature = "cargo-clippy", allow(too_many_arguments))]
    pub fn new<H>(
        shard_index: u64,
        shard_init: u64,
        shard_total: u64,
        ws_url: Arc<Mutex<String>>,
        token: Arc<Mutex<String>>,
        data: Arc<Mutex<ShareMap>>,
        event_handler: Arc<H>,
        framework: Arc<Mutex<Option<Box<Framework + Send>>>>,
        threadpool: ThreadPool,
    ) -> (Arc<Mutex<Self>>, ShardManagerMonitor) where H: EventHandler + Send + Sync + 'static {
        let (thread_tx, thread_rx) = mpsc::channel();
        let (shard_queue_tx, shard_queue_rx) = mpsc::channel();

        let runners = Arc::new(Mutex::new(HashMap::new()));

        let mut shard_queuer = ShardQueuer {
            data: Arc::clone(&data),
            event_handler: Arc::clone(&event_handler),
            framework: Arc::clone(&framework),
            last_start: None,
            manager_tx: thread_tx.clone(),
            runners: Arc::clone(&runners),
            rx: shard_queue_rx,
            token: Arc::clone(&token),
            ws_url: Arc::clone(&ws_url),
            threadpool,
        };

        thread::spawn(move || {
            shard_queuer.run();
        });

        let manager = Arc::new(Mutex::new(Self {
            shard_queuer: shard_queue_tx,
            runners,
            shard_index,
            shard_init,
            shard_total,
        }));

        (Arc::clone(&manager), ShardManagerMonitor {
            rx: thread_rx,
            manager,
        })
    }

    /// Creates a new shard manager, returning both the manager and a monitor
    /// for usage in a separate thread.
    #[cfg(not(feature = "framework"))]
    pub fn new<H>(
        shard_index: u64,
        shard_init: u64,
        shard_total: u64,
        ws_url: Arc<Mutex<String>>,
        token: Arc<Mutex<String>>,
        data: Arc<Mutex<ShareMap>>,
        event_handler: Arc<H>,
        threadpool: ThreadPool,
    ) -> (Arc<Mutex<Self>>, ShardManagerMonitor) where H: EventHandler + Send + Sync + 'static {
        let (thread_tx, thread_rx) = mpsc::channel();
        let (shard_queue_tx, shard_queue_rx) = mpsc::channel();

        let runners = Arc::new(Mutex::new(HashMap::new()));

        let mut shard_queuer = ShardQueuer {
            data: Arc::clone(&data),
            event_handler: Arc::clone(&event_handler),
            last_start: None,
            manager_tx: thread_tx.clone(),
            runners: Arc::clone(&runners),
            rx: shard_queue_rx,
            token: Arc::clone(&token),
            ws_url: Arc::clone(&ws_url),
            threadpool,
        };

        thread::spawn(move || {
            shard_queuer.run();
        });

        let manager = Arc::new(Mutex::new(Self {
            shard_queuer: shard_queue_tx,
            runners,
            shard_index,
            shard_init,
            shard_total,
        }));

        (Arc::clone(&manager), ShardManagerMonitor {
            rx: thread_rx,
            manager,
        })
    }

    /// Returns whether the shard manager contains either an active instance of
    /// a shard runner responsible for the given ID.
    ///
    /// If a shard has been queued but has not yet been initiated, then this
    /// will return `false`. Consider double-checking [`is_responsible_for`] to
    /// determine whether this shard manager is responsible for the given shard.
    ///
    /// [`is_responsible_for`]: #method.is_responsible_for
    pub fn has(&self, shard_id: ShardId) -> bool {
        self.runners.lock().contains_key(&shard_id)
    }

    /// Initializes all shards that the manager is responsible for.
    ///
    /// This will communicate shard boots with the [`ShardQueuer`] so that they
    /// are properly queued.
    ///
    /// [`ShardQueuer`]: struct.ShardQueuer.html
    pub fn initialize(&mut self) -> Result<()> {
        let shard_to = self.shard_index + self.shard_init;

        for shard_id in self.shard_index..shard_to {
            let shard_total = self.shard_total;

            self.boot([ShardId(shard_id), ShardId(shard_total)]);
        }

        Ok(())
    }

    /// Sets the new sharding information for the manager.
    ///
    /// This will shutdown all existing shards.
    ///
    /// This will _not_ instantiate the new shards.
    pub fn set_shards(&mut self, index: u64, init: u64, total: u64) {
        self.shutdown_all();

        self.shard_index = index;
        self.shard_init = init;
        self.shard_total = total;
    }

    /// Restarts a shard runner.
    ///
    /// This sends a shutdown signal to a shard's associated [`ShardRunner`],
    /// and then queues a initialization of a shard runner for the same shard
    /// via the [`ShardQueuer`].
    ///
    /// # Examples
    ///
    /// Creating a client and then restarting a shard by ID:
    ///
    /// _(note: in reality this precise code doesn't have an effect since the
    /// shard would not yet have been initialized via [`initialize`], but the
    /// concept is the same)_
    ///
    /// ```rust,no_run
    /// use serenity::client::bridge::gateway::ShardId;
    /// use serenity::client::{Client, EventHandler};
    /// use std::env;
    ///
    /// struct Handler;
    ///
    /// impl EventHandler for Handler { }
    ///
    /// let token = env::var("DISCORD_TOKEN").unwrap();
    /// let mut client = Client::new(&token, Handler).unwrap();
    ///
    /// // restart shard ID 7
    /// client.shard_manager.lock().restart(ShardId(7));
    /// ```
    ///
    /// [`ShardQueuer`]: struct.ShardQueuer.html
    /// [`ShardRunner`]: struct.ShardRunner.html
    pub fn restart(&mut self, shard_id: ShardId) {
        info!("Restarting shard {}", shard_id);
        self.shutdown(shard_id);

        let shard_total = self.shard_total;

        self.boot([shard_id, ShardId(shard_total)]);
    }

    /// Returns the [`ShardId`]s of the shards that have been instantiated and
    /// currently have a valid [`ShardRunner`].
    ///
    /// [`ShardId`]: struct.ShardId.html
    /// [`ShardRunner`]: struct.ShardRunner.html
    pub fn shards_instantiated(&self) -> Vec<ShardId> {
        self.runners.lock().keys().cloned().collect()
    }

    /// Attempts to shut down the shard runner by Id.
    ///
    /// Returns a boolean indicating whether a shard runner was present. This is
    /// _not_ necessary an indicator of whether the shard runner was
    /// successfully shut down.
    ///
    /// **Note**: If the receiving end of an mpsc channel - theoretically owned
    /// by the shard runner - no longer exists, then the shard runner will not
    /// know it should shut down. This _should never happen_. It may already be
    /// stopped.
    pub fn shutdown(&mut self, shard_id: ShardId) -> bool {
        info!("Shutting down shard {}", shard_id);

        if let Some(runner) = self.runners.lock().get(&shard_id) {
            let shutdown = ShardManagerMessage::Shutdown(shard_id);
            let msg = ShardClientMessage::Manager(shutdown);

            if let Err(why) = runner.runner_tx.send(msg) {
                warn!(
                    "Failed to cleanly shutdown shard {}: {:?}",
                    shard_id,
                    why,
                );
            }
        }

        self.runners.lock().remove(&shard_id).is_some()
    }

    /// Sends a shutdown message for all shards that the manager is responsible
    /// for that are still known to be running.
    ///
    /// If you only need to shutdown a select number of shards, prefer looping
    /// over the [`shutdown`] method.
    ///
    /// [`shutdown`]: #method.shutdown
    pub fn shutdown_all(&mut self) {
        let keys = {
            let runners = self.runners.lock();

            if runners.is_empty() {
                return;
            }

            runners.keys().cloned().collect::<Vec<_>>()
        };

        info!("Shutting down all shards");

        for shard_id in keys {
            self.shutdown(shard_id);
        }
    }

    fn boot(&mut self, shard_info: [ShardId; 2]) {
        info!("Telling shard queuer to start shard {}", shard_info[0]);

        let msg = ShardQueuerMessage::Start(shard_info[0], shard_info[1]);
        let _ = self.shard_queuer.send(msg);
    }
}

impl Drop for ShardManager {
    /// A custom drop implementation to clean up after the manager.
    ///
    /// This shuts down all active [`ShardRunner`]s and attempts to tell the
    /// [`ShardQueuer`] to shutdown.
    ///
    /// [`ShardQueuer`]: struct.ShardQueuer.html
    /// [`ShardRunner`]: struct.ShardRunner.html
    fn drop(&mut self) {
        self.shutdown_all();

        if let Err(why) = self.shard_queuer.send(ShardQueuerMessage::Shutdown) {
            warn!("Failed to send shutdown to shard queuer: {:?}", why);
        }
    }
}
<commit_msg>Make ShardManager::runners public<commit_after>use internal::prelude::*;
use parking_lot::Mutex;
use std::collections::HashMap;
use std::sync::mpsc::{self, Sender};
use std::sync::Arc;
use std::thread;
use super::super::super::EventHandler;
use super::{
    ShardClientMessage,
    ShardId,
    ShardManagerMessage,
    ShardManagerMonitor,
    ShardQueuer,
    ShardQueuerMessage,
    ShardRunnerInfo,
};
use threadpool::ThreadPool;
use typemap::ShareMap;

#[cfg(feature = "framework")]
use framework::Framework;

/// A manager for handling the status of shards by starting them, restarting
/// them, and stopping them when required.
///
/// **Note**: The [`Client`] internally uses a shard manager. If you are using a
/// Client, then you do not need to make one of these.
///
/// # Examples
///
/// Initialize a shard manager with a framework responsible for shards 0 through
/// 2, of 5 total shards:
///
/// ```rust,no_run
/// extern crate parking_lot;
/// extern crate serenity;
/// extern crate threadpool;
/// extern crate typemap;
///
/// # use std::error::Error;
/// #
/// # #[cfg(feature = "framework")]
/// # fn try_main() -> Result<(), Box<Error>> {
/// #
/// use parking_lot::Mutex;
/// use serenity::client::bridge::gateway::ShardManager;
/// use serenity::client::EventHandler;
/// use serenity::http;
/// use std::sync::Arc;
/// use std::env;
/// use threadpool::ThreadPool;
/// use typemap::ShareMap;
///
/// struct Handler;
///
/// impl EventHandler for Handler { }
///
/// let token = env::var("DISCORD_TOKEN")?;
/// http::set_token(&token);
/// let token = Arc::new(Mutex::new(token));
///
/// let gateway_url = Arc::new(Mutex::new(http::get_gateway()?.url));
/// let data = Arc::new(Mutex::new(ShareMap::custom()));
/// let event_handler = Arc::new(Handler);
/// let framework = Arc::new(Mutex::new(None));
/// let threadpool = ThreadPool::with_name("my threadpool".to_owned(), 5);
///
/// ShardManager::new(
///     0, // the shard index to start initiating from
///     3, // the number of shards to initiate (this initiates 0, 1, and 2)
///     5, // the total number of shards in use
///     gateway_url,
///     token,
///     data,
///     event_handler,
///     framework,
///     threadpool,
/// );
/// #     Ok(())
/// # }
/// #
/// # #[cfg(not(feature = "framework"))]
/// # fn try_main() -> Result<(), Box<Error>> {
/// #     Ok(())
/// # }
/// #
/// # fn main() {
/// #     try_main().unwrap();
/// # }
/// ```
///
/// [`Client`]: ../../struct.Client.html
#[derive(Debug)]
pub struct ShardManager {
    /// The shard runners currently managed.
    ///
    /// **Note**: It is highly unrecommended to mutate this yourself unless you
    /// need to. Instead prefer to use methods on this struct that are provided
    /// where possible.
    pub runners: Arc<Mutex<HashMap<ShardId, ShardRunnerInfo>>>,
    /// The index of the first shard to initialize, 0-indexed.
    shard_index: u64,
    /// The number of shards to initialize.
    shard_init: u64,
    /// The total shards in use, 1-indexed.
    shard_total: u64,
    shard_queuer: Sender<ShardQueuerMessage>,
}

impl ShardManager {
    /// Creates a new shard manager, returning both the manager and a monitor
    /// for usage in a separate thread.
    #[cfg(feature = "framework")]
    #[cfg_attr(feature = "cargo-clippy", allow(too_many_arguments))]
    pub fn new<H>(
        shard_index: u64,
        shard_init: u64,
        shard_total: u64,
        ws_url: Arc<Mutex<String>>,
        token: Arc<Mutex<String>>,
        data: Arc<Mutex<ShareMap>>,
        event_handler: Arc<H>,
        framework: Arc<Mutex<Option<Box<Framework + Send>>>>,
        threadpool: ThreadPool,
    ) -> (Arc<Mutex<Self>>, ShardManagerMonitor) where H: EventHandler + Send + Sync + 'static {
        let (thread_tx, thread_rx) = mpsc::channel();
        let (shard_queue_tx, shard_queue_rx) = mpsc::channel();

        let runners = Arc::new(Mutex::new(HashMap::new()));

        let mut shard_queuer = ShardQueuer {
            data: Arc::clone(&data),
            event_handler: Arc::clone(&event_handler),
            framework: Arc::clone(&framework),
            last_start: None,
            manager_tx: thread_tx.clone(),
            runners: Arc::clone(&runners),
            rx: shard_queue_rx,
            token: Arc::clone(&token),
            ws_url: Arc::clone(&ws_url),
            threadpool,
        };

        thread::spawn(move || {
            shard_queuer.run();
        });

        let manager = Arc::new(Mutex::new(Self {
            shard_queuer: shard_queue_tx,
            runners,
            shard_index,
            shard_init,
            shard_total,
        }));

        (Arc::clone(&manager), ShardManagerMonitor {
            rx: thread_rx,
            manager,
        })
    }

    /// Creates a new shard manager, returning both the manager and a monitor
    /// for usage in a separate thread.
    #[cfg(not(feature = "framework"))]
    pub fn new<H>(
        shard_index: u64,
        shard_init: u64,
        shard_total: u64,
        ws_url: Arc<Mutex<String>>,
        token: Arc<Mutex<String>>,
        data: Arc<Mutex<ShareMap>>,
        event_handler: Arc<H>,
        threadpool: ThreadPool,
    ) -> (Arc<Mutex<Self>>, ShardManagerMonitor) where H: EventHandler + Send + Sync + 'static {
        let (thread_tx, thread_rx) = mpsc::channel();
        let (shard_queue_tx, shard_queue_rx) = mpsc::channel();

        let runners = Arc::new(Mutex::new(HashMap::new()));

        let mut shard_queuer = ShardQueuer {
            data: Arc::clone(&data),
            event_handler: Arc::clone(&event_handler),
            last_start: None,
            manager_tx: thread_tx.clone(),
            runners: Arc::clone(&runners),
            rx: shard_queue_rx,
            token: Arc::clone(&token),
            ws_url: Arc::clone(&ws_url),
            threadpool,
        };

        thread::spawn(move || {
            shard_queuer.run();
        });

        let manager = Arc::new(Mutex::new(Self {
            shard_queuer: shard_queue_tx,
            runners,
            shard_index,
            shard_init,
            shard_total,
        }));

        (Arc::clone(&manager), ShardManagerMonitor {
            rx: thread_rx,
            manager,
        })
    }

    /// Returns whether the shard manager contains either an active instance of
    /// a shard runner responsible for the given ID.
    ///
    /// If a shard has been queued but has not yet been initiated, then this
    /// will return `false`. Consider double-checking [`is_responsible_for`] to
    /// determine whether this shard manager is responsible for the given shard.
    ///
    /// [`is_responsible_for`]: #method.is_responsible_for
    pub fn has(&self, shard_id: ShardId) -> bool {
        self.runners.lock().contains_key(&shard_id)
    }

    /// Initializes all shards that the manager is responsible for.
    ///
    /// This will communicate shard boots with the [`ShardQueuer`] so that they
    /// are properly queued.
    ///
    /// [`ShardQueuer`]: struct.ShardQueuer.html
    pub fn initialize(&mut self) -> Result<()> {
        let shard_to = self.shard_index + self.shard_init;

        for shard_id in self.shard_index..shard_to {
            let shard_total = self.shard_total;

            self.boot([ShardId(shard_id), ShardId(shard_total)]);
        }

        Ok(())
    }

    /// Sets the new sharding information for the manager.
    ///
    /// This will shutdown all existing shards.
    ///
    /// This will _not_ instantiate the new shards.
    pub fn set_shards(&mut self, index: u64, init: u64, total: u64) {
        self.shutdown_all();

        self.shard_index = index;
        self.shard_init = init;
        self.shard_total = total;
    }

    /// Restarts a shard runner.
    ///
    /// This sends a shutdown signal to a shard's associated [`ShardRunner`],
    /// and then queues a initialization of a shard runner for the same shard
    /// via the [`ShardQueuer`].
    ///
    /// # Examples
    ///
    /// Creating a client and then restarting a shard by ID:
    ///
    /// _(note: in reality this precise code doesn't have an effect since the
    /// shard would not yet have been initialized via [`initialize`], but the
    /// concept is the same)_
    ///
    /// ```rust,no_run
    /// use serenity::client::bridge::gateway::ShardId;
    /// use serenity::client::{Client, EventHandler};
    /// use std::env;
    ///
    /// struct Handler;
    ///
    /// impl EventHandler for Handler { }
    ///
    /// let token = env::var("DISCORD_TOKEN").unwrap();
    /// let mut client = Client::new(&token, Handler).unwrap();
    ///
    /// // restart shard ID 7
    /// client.shard_manager.lock().restart(ShardId(7));
    /// ```
    ///
    /// [`ShardQueuer`]: struct.ShardQueuer.html
    /// [`ShardRunner`]: struct.ShardRunner.html
    pub fn restart(&mut self, shard_id: ShardId) {
        info!("Restarting shard {}", shard_id);
        self.shutdown(shard_id);

        let shard_total = self.shard_total;

        self.boot([shard_id, ShardId(shard_total)]);
    }

    /// Returns the [`ShardId`]s of the shards that have been instantiated and
    /// currently have a valid [`ShardRunner`].
    ///
    /// [`ShardId`]: struct.ShardId.html
    /// [`ShardRunner`]: struct.ShardRunner.html
    pub fn shards_instantiated(&self) -> Vec<ShardId> {
        self.runners.lock().keys().cloned().collect()
    }

    /// Attempts to shut down the shard runner by Id.
    ///
    /// Returns a boolean indicating whether a shard runner was present. This is
    /// _not_ necessary an indicator of whether the shard runner was
    /// successfully shut down.
    ///
    /// **Note**: If the receiving end of an mpsc channel - theoretically owned
    /// by the shard runner - no longer exists, then the shard runner will not
    /// know it should shut down. This _should never happen_. It may already be
    /// stopped.
    pub fn shutdown(&mut self, shard_id: ShardId) -> bool {
        info!("Shutting down shard {}", shard_id);

        if let Some(runner) = self.runners.lock().get(&shard_id) {
            let shutdown = ShardManagerMessage::Shutdown(shard_id);
            let msg = ShardClientMessage::Manager(shutdown);

            if let Err(why) = runner.runner_tx.send(msg) {
                warn!(
                    "Failed to cleanly shutdown shard {}: {:?}",
                    shard_id,
                    why,
                );
            }
        }

        self.runners.lock().remove(&shard_id).is_some()
    }

    /// Sends a shutdown message for all shards that the manager is responsible
    /// for that are still known to be running.
    ///
    /// If you only need to shutdown a select number of shards, prefer looping
    /// over the [`shutdown`] method.
    ///
    /// [`shutdown`]: #method.shutdown
    pub fn shutdown_all(&mut self) {
        let keys = {
            let runners = self.runners.lock();

            if runners.is_empty() {
                return;
            }

            runners.keys().cloned().collect::<Vec<_>>()
        };

        info!("Shutting down all shards");

        for shard_id in keys {
            self.shutdown(shard_id);
        }
    }

    fn boot(&mut self, shard_info: [ShardId; 2]) {
        info!("Telling shard queuer to start shard {}", shard_info[0]);

        let msg = ShardQueuerMessage::Start(shard_info[0], shard_info[1]);
        let _ = self.shard_queuer.send(msg);
    }
}

impl Drop for ShardManager {
    /// A custom drop implementation to clean up after the manager.
    ///
    /// This shuts down all active [`ShardRunner`]s and attempts to tell the
    /// [`ShardQueuer`] to shutdown.
    ///
    /// [`ShardQueuer`]: struct.ShardQueuer.html
    /// [`ShardRunner`]: struct.ShardRunner.html
    fn drop(&mut self) {
        self.shutdown_all();

        if let Err(why) = self.shard_queuer.send(ShardQueuerMessage::Shutdown) {
            warn!("Failed to send shutdown to shard queuer: {:?}", why);
        }
    }
}
<|endoftext|>


<commit_before>#[derive(PartialEq, Debug)]
enum MqttType {
    Connect,
    Ping,
}

fn message_type(bytes: &[u8]) -> MqttType {
    MqttType::Connect
}

#[test]
fn connect_type() {
    let connect_bytes = &[
        0x10u8, 0x2a, // fixed header
        0x00, 0x06, 'M' as u8, 'Q' as u8, 'I' as u8, 's' as u8, 'd' as u8, 'p' as u8,
        0x03, // protocol version
        0xcc, // connection flags 1100111x user, pw, !wr, w(01), w, !c, x
        0x00, 0x0a, // keepalive of 100
        0x00, 0x03, 'c' as u8, 'i' as u8, 'd' as u8, // client ID
        0x00, 0x04, 'w' as u8, 'i' as u8, 'l' as u8, 'l' as u8, // will topic
        0x00, 0x04, 'w' as u8, 'm' as u8, 's' as u8, 'g' as u8, // will msg
        0x00, 0x07, 'g' as u8, 'l' as u8, 'i' as u8, 'f' as u8, 't' as u8, 'e' as u8, 'l' as u8, // username
        0x00, 0x02, 'p' as u8, 'w' as u8, // password
        ][0..];
    assert_eq!(message_type(connect_bytes), MqttType::Connect);
}

#[test]
fn ping_type() {
    let ping_bytes = &[0xc0u8, 0][0..];
    assert_eq!(message_type(ping_bytes), MqttType::Ping);
}
<commit_msg>Unit test passes for message type<commit_after>use std::mem;

#[derive(PartialEq, Debug)]
enum MqttType {
    Reserved = 0,
    Connect,
    ConnAck,
    Publish,
    PubAck,
    PubRec,
    PubRel,
    PubComp,
    Subscribe,
    SubAck,
    Unsubscribe,
    UnsubAck,
    PingReq,
    PingResp,
    Disconnect,
}

fn message_type(bytes: &[u8]) -> MqttType {
    unsafe { mem::transmute((bytes[0] & 0xf0) >> 4) }
}

#[test]
fn connect_type() {
    let connect_bytes = &[
        0x10u8, 0x2a, // fixed header
        0x00, 0x06, 'M' as u8, 'Q' as u8, 'I' as u8, 's' as u8, 'd' as u8, 'p' as u8,
        0x03, // protocol version
        0xcc, // connection flags 1100111x user, pw, !wr, w(01), w, !c, x
        0x00, 0x0a, // keepalive of 100
        0x00, 0x03, 'c' as u8, 'i' as u8, 'd' as u8, // client ID
        0x00, 0x04, 'w' as u8, 'i' as u8, 'l' as u8, 'l' as u8, // will topic
        0x00, 0x04, 'w' as u8, 'm' as u8, 's' as u8, 'g' as u8, // will msg
        0x00, 0x07, 'g' as u8, 'l' as u8, 'i' as u8, 'f' as u8, 't' as u8, 'e' as u8, 'l' as u8, // username
        0x00, 0x02, 'p' as u8, 'w' as u8, // password
        ][0..];
    assert_eq!(message_type(connect_bytes), MqttType::Connect);
}

#[test]
fn ping_type() {
    let ping_bytes = &[0xc0u8, 0][0..];
    assert_eq!(message_type(ping_bytes), MqttType::PingReq);
}
<|endoftext|>


<commit_before>use super::*;

pub struct CreateDocumentRequest {
    hyper_client: HyperClient,
    request: RequestBuilder,
    payload: Result<String, serde_json::Error>,
}

impl DocumentRequestExt for CreateDocumentRequest {
    fn request(&mut self) -> &mut RequestBuilder {
        &mut self.request
    }
}

impl CreateDocumentRequest {
    pub(crate) fn new(
        hyper_client: HyperClient,
        request: RequestBuilder,
        payload: Result<String, serde_json::Error>,
    ) -> CreateDocumentRequest {
        CreateDocumentRequest {
            hyper_client,
            request,
            payload,
        }
    }

    request_option!(is_upsert, bool, HEADER_DOCUMENTDB_IS_UPSERT);
    request_option!(indexing_directive, IndexingDirective, HEADER_INDEXING_DIRECTIVE);

    pub fn execute(self) -> impl Future<Item = DocumentAttributes, Error = AzureError> {
        trace!("get_document called(request == {:?}", self.request);
        let hc = self.hyper_client;
        let mut req = self.request;
        future::result(self.payload)
            .from_err()
            .and_then(move |payload| Ok(req.body(payload.into())?))
            .and_then(move |r| check_status_extract_body(hc.request(r), StatusCode::CREATED))
            .and_then(move |body| Ok(serde_json::from_str::<DocumentAttributes>(&body)?))
    }
}

pub struct GetDocumentRequest {
    hyper_client: HyperClient,
    request: RequestBuilder,
}

impl DocumentRequestExt for GetDocumentRequest {
    fn request(&mut self) -> &mut RequestBuilder {
        &mut self.request
    }
}

impl GetDocumentRequest {
    pub(crate) fn new(hyper_client: HyperClient, request: RequestBuilder) -> GetDocumentRequest {
        GetDocumentRequest { hyper_client, request }
    }

    request_bytes_option!(if_none_match, String, header::IF_NONE_MATCH);

    pub fn execute<T: DeserializeOwned>(mut self) -> impl Future<Item = GetDocumentResponse<T>, Error = AzureError> {
        trace!("get_document called(request == {:?}", self.request);

        future::result(self.request.body(hyper::Body::empty()))
            .from_err()
            .and_then(move |r| extract_status_headers_and_body(self.hyper_client.request(r)))
            .and_then(move |(status, headers, body)| Self::extract_result(status, &headers, &body))
    }

    fn extract_result<R: DeserializeOwned>(
        status: hyper::StatusCode,
        headers: &HeaderMap,
        body: &[u8],
    ) -> Result<GetDocumentResponse<R>, AzureError> {
        match status {
            StatusCode::OK => {
                let additional_headers = DocumentAdditionalHeaders::derive_from(headers);
                let document = Document::from_json(body)?;
                Ok(GetDocumentResponse {
                    document: Some(document),
                    additional_headers,
                })
            }
            // NotFound is not an error so we return None along
            // with the additional headers.
            StatusCode::NOT_FOUND => {
                let additional_headers = DocumentAdditionalHeaders::derive_from(headers);
                Ok(GetDocumentResponse {
                    document: None,
                    additional_headers,
                })
            }
            _ => {
                // We treat everything else as an error. We could
                // handle 304 (Not modified) in a specific way but
                // for now we do not.
                let error_text = str::from_utf8(body)?;
                Err(AzureError::UnexpectedHTTPResult(UnexpectedHTTPResult::new(
                    StatusCode::OK,
                    status,
                    error_text,
                )))
            }
        }
    }
}

pub struct QueryDocumentRequest {
    hyper_client: HyperClient,
    request: RequestBuilder,
    payload: Result<String, serde_json::Error>,
}

impl DocumentRequestExt for QueryDocumentRequest {
    fn request(&mut self) -> &mut RequestBuilder {
        &mut self.request
    }
}

const QUERY_CONTENT_TYPE: &str = "application/query+json";

impl QueryDocumentRequest {
    pub(crate) fn new(
        hyper_client: HyperClient,
        mut request: RequestBuilder,
        payload: Result<String, serde_json::Error>,
    ) -> QueryDocumentRequest {
        request
            .header(HEADER_DOCUMENTDB_ISQUERY, HeaderValue::from_static("true"))
            .header(header::CONTENT_TYPE, HeaderValue::from_static(QUERY_CONTENT_TYPE));
        QueryDocumentRequest {
            hyper_client,
            request,
            payload,
        }
    }

    request_option!(max_item_count, u64, HEADER_MAX_ITEM_COUNT);
    request_bytes_option!(continuation_token, ContinuationToken, HEADER_CONTINUATION);
    request_option!(enable_cross_partition, bool, HEADER_DOCUMENTDB_QUERY_ENABLECROSSPARTITION);
    request_option!(consistency_level, ConsistencyLevel, HEADER_CONSISTENCY_LEVEL);

    pub fn execute<T: DeserializeOwned>(self) -> impl Future<Item = QueryDocumentResponse<T>, Error = AzureError> {
        trace!("get_document called(request == {:?}", self.request);
        self.execute_json().and_then(Self::convert_query_document_type)
    }

    pub fn execute_json(self) -> impl Future<Item = QueryDocumentResponse<serde_json::Value>, Error = AzureError> {
        trace!("query_document called(request == {:?}", self.request);
        let hc = self.hyper_client;
        let mut req = self.request;
        future::result(self.payload)
            .from_err()
            .and_then(move |payload| Ok(req.body(payload.into())?))
            .and_then(move |r| check_status_extract_headers_and_body(hc.request(r), StatusCode::OK))
            .and_then(move |(headers, body)| Self::extract_result_json(&body, &headers))
    }

    fn extract_result_json(body: &[u8], headers: &HeaderMap) -> Result<QueryDocumentResponse<serde_json::Value>, AzureError> {
        trace!("headers == {:?}", headers);

        let additional_headers = QueryDocumentResponseAdditonalHeaders {
            // This match just tries to extract the info and convert it
            // into the correct type. It is complicated because headers
            // can be missing and also because headers.get<T> will return
            // a T reference (&T) so we need to cast it into the
            // correct type and clone it (in this case into a &str that will
            // become a String using to_owned())
            continuation_token: derive_continuation_token(headers),
            // Here we assume the Charge header to always be present.
            // If problems arise we
            // will change the field to be Option(al).
            charge: derive_request_charge(headers),
        };
        debug!("additional_headers == {:?}", additional_headers);

        let query_response_meta = serde_json::from_slice::<QueryResponseMeta>(body)?;
        debug!("query_response_meta == {:?}", &query_response_meta);

        let json = str::from_utf8(body)?;
        debug!("json == {}", json);

        let mut v: serde_json::Value = serde_json::from_slice(body)?;

        // Work on Documents section
        let mut d = v.get_mut("Documents").unwrap().take();
        debug!("\n\nd == {:?}\n\n", d);

        let docs = d.as_array_mut().unwrap().into_iter().map(|doc| {
            // We could either have a Document or a plain entry.
            // We will find out here.
            let mut doc = doc.take();

            let attrs = {
                let map = doc.as_object_mut().unwrap();
                DocumentAttributes::try_extract(map)
            };

            // debug!("\ndocument_attributes == {:?}", document_attributes);

            QueryResult {
                document_attributes: attrs,
                result: doc,
            }
        });

        Ok(QueryDocumentResponse {
            query_response_meta,
            additional_headers,
            results: docs.collect(),
        })
    }

    #[inline]
    fn convert_query_document_type<T>(qdr: QueryDocumentResponse<serde_json::Value>) -> Result<QueryDocumentResponse<T>, AzureError>
    where
        T: DeserializeOwned,
    {
        let mut qdr_converted: QueryDocumentResponse<T> = QueryDocumentResponse {
            query_response_meta: qdr.query_response_meta,
            results: Vec::new(),
            additional_headers: qdr.additional_headers,
        };

        for res_json in qdr.results {
            qdr_converted.results.push(QueryResult {
                document_attributes: res_json.document_attributes,
                result: serde_json::from_value(res_json.result)?,
            });
        }

        Ok(qdr_converted)
    }
}

pub struct ListDocumentsRequest {
    hyper_client: HyperClient,
    request: RequestBuilder,
}

impl ListDocumentsRequest {
    pub(crate) fn new(hyper_client: HyperClient, request: RequestBuilder) -> ListDocumentsRequest {
        ListDocumentsRequest { hyper_client, request }
    }

    request_option!(max_item_count, u64, HEADER_MAX_ITEM_COUNT);
    request_bytes_option!(continuation_token, ContinuationToken, HEADER_CONTINUATION);
    request_option!(consistency_level, ConsistencyLevel, HEADER_CONSISTENCY_LEVEL);
    request_bytes_option!(session_token, String, HEADER_SESSION_TOKEN);
    request_bytes_option!(if_none_match, String, header::IF_NONE_MATCH);
    request_bytes_option!(partition_range_id, String, HEADER_DOCUMENTDB_PARTITIONRANGEID);

    pub fn incremental_feed(mut self) -> Self {
        self.request.header(HEADER_A_IM, HeaderValue::from_static("Incremental feed"));
        self
    }

    pub fn execute<T: DeserializeOwned>(mut self) -> impl Future<Item = ListDocumentsResponse<T>, Error = AzureError> {
        future::result(self.request.body(hyper::Body::empty()))
            .from_err()
            .and_then(move |r| check_status_extract_headers_and_body(self.hyper_client.request(r), StatusCode::OK))
            .and_then(|(headers, whole_body)| Self::extract_result::<T>(&whole_body, &headers))
    }

    fn extract_result<T>(body: &[u8], headers: &HeaderMap) -> Result<ListDocumentsResponse<T>, AzureError>
    where
        T: DeserializeOwned,
    {
        debug!("headers == {:?}", headers);

        let ado = ListDocumentsResponseAdditionalHeaders {
            // This match just tries to extract the info and convert it
            // into the correct type. It is complicated because headers
            // can be missing and also because headers.get<T> will return
            // a T reference (&T) so we need to cast it into the
            // correct type and clone it (in this case into a &str that will
            // become a String using to_owned())
            continuation_token: derive_continuation_token(headers),
            // Here we assume the Charge header to always be present.
            // If problems arise we
            // will change the field to be Option(al).
            charge: derive_request_charge(headers),
            etag: headers.get(header::ETAG).and_then(|v| v.to_str().ok()).map(|s| s.to_owned()),
        };
        debug!("ado == {:?}", ado);

        // we will proceed in three steps:
        // 1- Deserialize the result as DocumentAttributes. The extra field will be ignored.
        // 2- Deserialize the result a type T. The extra fields will be ignored.
        // 3- Zip 1 and 2 in the resulting structure.
        // There is a lot of data movement here, let's hope the compiler is smarter than me :)
        let document_attributes = serde_json::from_slice::<ListDocumentsResponseAttributes>(body)?;
        let entries = serde_json::from_slice::<ListDocumentsResponseEntities<T>>(body)?;

        let documents = document_attributes
            .documents
            .into_iter()
            .zip(entries.entities.into_iter())
            .map(|(da, e)| Document {
                document_attributes: da,
                entity: e,
            })
            .collect();

        Ok(ListDocumentsResponse {
            rid: document_attributes.rid,
            documents,
            additional_headers: ado,
        })
    }
}

pub struct ReplaceDocumentRequest<T> {
    hyper_client: HyperClient,
    request: RequestBuilder,
    payload: Result<String, serde_json::Error>,
    _t: PhantomData<T>,
}

impl<T> DocumentRequestExt for ReplaceDocumentRequest<T> {
    fn request(&mut self) -> &mut RequestBuilder {
        &mut self.request
    }
}

impl<T: DeserializeOwned> ReplaceDocumentRequest<T> {
    pub(crate) fn new(
        hyper_client: HyperClient,
        request: RequestBuilder,
        payload: Result<String, serde_json::Error>,
    ) -> ReplaceDocumentRequest<T> {
        ReplaceDocumentRequest {
            hyper_client,
            request,
            payload,
            _t: PhantomData,
        }
    }

    request_bytes_option!(if_match, String, header::IF_MATCH);
    request_option!(indexing_directive, IndexingDirective, HEADER_INDEXING_DIRECTIVE);

    pub fn execute(self) -> impl Future<Item = ReplaceDocumentResponse<T>, Error = AzureError> {
        trace!("get_document called(request == {:?}", self.request);
        let hc = self.hyper_client;
        let mut req = self.request;
        future::result(self.payload)
            .from_err()
            .and_then(move |payload| Ok(req.body(payload.into())?))
            .and_then(move |r| check_status_extract_headers_and_body(hc.request(r), StatusCode::OK))
            .and_then(move |(headers, body)| Self::extract_result(&headers, &body))
    }

    fn extract_result<R: DeserializeOwned>(headers: &HeaderMap, body: &[u8]) -> Result<ReplaceDocumentResponse<R>, AzureError> {
        let additional_headers = DocumentAdditionalHeaders::derive_from(headers);
        let document = Document::from_json(body)?;
        Ok(ReplaceDocumentResponse {
            document,
            additional_headers,
        })
    }
}

pub struct DeleteDocumentRequest {
    hyper_client: HyperClient,
    request: RequestBuilder,
}

impl DocumentRequestExt for DeleteDocumentRequest {
    fn request(&mut self) -> &mut RequestBuilder {
        &mut self.request
    }
}

impl DeleteDocumentRequest {
    pub(crate) fn new(hyper_client: HyperClient, request: RequestBuilder) -> DeleteDocumentRequest {
        DeleteDocumentRequest { hyper_client, request }
    }

    request_bytes_option!(if_match, String, header::IF_MATCH);

    pub fn execute(mut self) -> impl Future<Item = (), Error = AzureError> {
        trace!("get_document called(request == {:?}", self.request);

        future::result(self.request.body(hyper::Body::empty()))
            .from_err()
            .and_then(move |r| check_status_extract_body(self.hyper_client.request(r), StatusCode::NO_CONTENT))
            .and_then(|_| Ok(()))
    }
}

pub trait DocumentRequestExt: Sized {
    fn request(&mut self) -> &mut RequestBuilder;

    fn session_token<S: AsRef<str>>(mut self, token: S) -> Self {
        self.request().header_formatted(HEADER_SESSION_TOKEN, token.as_ref());
        self
    }

    fn partition_key<'a, P: Into<PartitionKey<'a>>>(mut self, key: P) -> Self {
        // todo: move unwrap into PartitionKey impl itself as we control the impl and it surely won't error out
        if let Some(ser_key) = key.into().to_json().unwrap() {
            self.request().header_formatted(HEADER_DOCUMENTDB_PARTITIONKEY, ser_key);
        }
        self
    }
}

fn derive_continuation_token(headers: &HeaderMap) -> Option<String> {
    headers.get(HEADER_CONTINUATION).and_then(|v| v.to_str().ok()).map(|v| v.to_owned())
}

fn derive_request_charge(headers: &HeaderMap) -> f64 {
    headers.get(HEADER_REQUEST_CHARGE).unwrap().to_str().unwrap().parse().unwrap()
}
<commit_msg>Fixed cosmos queries not returning DocumentAttributes json<commit_after>use super::*;

pub struct CreateDocumentRequest {
    hyper_client: HyperClient,
    request: RequestBuilder,
    payload: Result<String, serde_json::Error>,
}

impl DocumentRequestExt for CreateDocumentRequest {
    fn request(&mut self) -> &mut RequestBuilder {
        &mut self.request
    }
}

impl CreateDocumentRequest {
    pub(crate) fn new(
        hyper_client: HyperClient,
        request: RequestBuilder,
        payload: Result<String, serde_json::Error>,
    ) -> CreateDocumentRequest {
        CreateDocumentRequest {
            hyper_client,
            request,
            payload,
        }
    }

    request_option!(is_upsert, bool, HEADER_DOCUMENTDB_IS_UPSERT);
    request_option!(indexing_directive, IndexingDirective, HEADER_INDEXING_DIRECTIVE);

    pub fn execute(self) -> impl Future<Item = DocumentAttributes, Error = AzureError> {
        trace!("get_document called(request == {:?}", self.request);
        let hc = self.hyper_client;
        let mut req = self.request;
        future::result(self.payload)
            .from_err()
            .and_then(move |payload| Ok(req.body(payload.into())?))
            .and_then(move |r| check_status_extract_body(hc.request(r), StatusCode::CREATED))
            .and_then(move |body| Ok(serde_json::from_str::<DocumentAttributes>(&body)?))
    }
}

pub struct GetDocumentRequest {
    hyper_client: HyperClient,
    request: RequestBuilder,
}

impl DocumentRequestExt for GetDocumentRequest {
    fn request(&mut self) -> &mut RequestBuilder {
        &mut self.request
    }
}

impl GetDocumentRequest {
    pub(crate) fn new(hyper_client: HyperClient, request: RequestBuilder) -> GetDocumentRequest {
        GetDocumentRequest { hyper_client, request }
    }

    request_bytes_option!(if_none_match, String, header::IF_NONE_MATCH);

    pub fn execute<T: DeserializeOwned>(mut self) -> impl Future<Item = GetDocumentResponse<T>, Error = AzureError> {
        trace!("get_document called(request == {:?}", self.request);

        future::result(self.request.body(hyper::Body::empty()))
            .from_err()
            .and_then(move |r| extract_status_headers_and_body(self.hyper_client.request(r)))
            .and_then(move |(status, headers, body)| Self::extract_result(status, &headers, &body))
    }

    fn extract_result<R: DeserializeOwned>(
        status: hyper::StatusCode,
        headers: &HeaderMap,
        body: &[u8],
    ) -> Result<GetDocumentResponse<R>, AzureError> {
        match status {
            StatusCode::OK => {
                let additional_headers = DocumentAdditionalHeaders::derive_from(headers);
                let document = Document::from_json(body)?;
                Ok(GetDocumentResponse {
                    document: Some(document),
                    additional_headers,
                })
            }
            // NotFound is not an error so we return None along
            // with the additional headers.
            StatusCode::NOT_FOUND => {
                let additional_headers = DocumentAdditionalHeaders::derive_from(headers);
                Ok(GetDocumentResponse {
                    document: None,
                    additional_headers,
                })
            }
            _ => {
                // We treat everything else as an error. We could
                // handle 304 (Not modified) in a specific way but
                // for now we do not.
                let error_text = str::from_utf8(body)?;
                Err(AzureError::UnexpectedHTTPResult(UnexpectedHTTPResult::new(
                    StatusCode::OK,
                    status,
                    error_text,
                )))
            }
        }
    }
}

pub struct QueryDocumentRequest {
    hyper_client: HyperClient,
    request: RequestBuilder,
    payload: Result<String, serde_json::Error>,
}

impl DocumentRequestExt for QueryDocumentRequest {
    fn request(&mut self) -> &mut RequestBuilder {
        &mut self.request
    }
}

const QUERY_CONTENT_TYPE: &str = "application/query+json";

impl QueryDocumentRequest {
    pub(crate) fn new(
        hyper_client: HyperClient,
        mut request: RequestBuilder,
        payload: Result<String, serde_json::Error>,
    ) -> QueryDocumentRequest {
        request
            .header(HEADER_DOCUMENTDB_ISQUERY, HeaderValue::from_static("true"))
            .header(header::CONTENT_TYPE, HeaderValue::from_static(QUERY_CONTENT_TYPE));
        QueryDocumentRequest {
            hyper_client,
            request,
            payload,
        }
    }

    request_option!(max_item_count, u64, HEADER_MAX_ITEM_COUNT);
    request_bytes_option!(continuation_token, ContinuationToken, HEADER_CONTINUATION);
    request_option!(enable_cross_partition, bool, HEADER_DOCUMENTDB_QUERY_ENABLECROSSPARTITION);
    request_option!(consistency_level, ConsistencyLevel, HEADER_CONSISTENCY_LEVEL);

    pub fn execute<T: DeserializeOwned>(self) -> impl Future<Item = QueryDocumentResponse<T>, Error = AzureError> {
        trace!("get_document called(request == {:?}", self.request);
        self.execute_json().and_then(Self::convert_query_document_type)
    }

    pub fn execute_json(self) -> impl Future<Item = QueryDocumentResponse<serde_json::Value>, Error = AzureError> {
        trace!("query_document called(request == {:?}", self.request);
        let hc = self.hyper_client;
        let mut req = self.request;
        future::result(self.payload)
            .from_err()
            .and_then(move |payload| Ok(req.body(payload.into())?))
            .and_then(move |r| check_status_extract_headers_and_body(hc.request(r), StatusCode::OK))
            .and_then(move |(headers, body)| Self::extract_result_json(&body, &headers))
    }

    fn extract_result_json(body: &[u8], headers: &HeaderMap) -> Result<QueryDocumentResponse<serde_json::Value>, AzureError> {
        trace!("headers == {:?}", headers);

        let additional_headers = QueryDocumentResponseAdditonalHeaders {
            // This match just tries to extract the info and convert it
            // into the correct type. It is complicated because headers
            // can be missing and also because headers.get<T> will return
            // a T reference (&T) so we need to cast it into the
            // correct type and clone it (in this case into a &str that will
            // become a String using to_owned())
            continuation_token: derive_continuation_token(headers),
            // Here we assume the Charge header to always be present.
            // If problems arise we
            // will change the field to be Option(al).
            charge: derive_request_charge(headers),
        };
        debug!("additional_headers == {:?}", additional_headers);

        let query_response_meta = serde_json::from_slice::<QueryResponseMeta>(body)?;
        debug!("query_response_meta == {:?}", &query_response_meta);

        let json = str::from_utf8(body)?;
        debug!("json == {}", json);

        let mut v: serde_json::Value = serde_json::from_slice(body)?;

        // Work on Documents section
        let mut d = v.get_mut("Documents").unwrap().take();
        debug!("\n\nd == {:?}\n\n", d);

        let docs = d.as_array_mut().unwrap().into_iter().map(|doc| {
            // We could either have a Document or a plain entry.
            // We will find out here.
            let mut doc = doc.take();

            let attrs = {
                if let Some(map) = doc.as_object_mut() {
                    DocumentAttributes::try_extract(map)
                } else {
                    None
                }
            };

            debug!("attrs == {:?}", attrs);

            QueryResult {
                document_attributes: attrs,
                result: doc,
            }
        });

        Ok(QueryDocumentResponse {
            query_response_meta,
            additional_headers,
            results: docs.collect(),
        })
    }

    #[inline]
    fn convert_query_document_type<T>(qdr: QueryDocumentResponse<serde_json::Value>) -> Result<QueryDocumentResponse<T>, AzureError>
    where
        T: DeserializeOwned,
    {
        let mut qdr_converted: QueryDocumentResponse<T> = QueryDocumentResponse {
            query_response_meta: qdr.query_response_meta,
            results: Vec::new(),
            additional_headers: qdr.additional_headers,
        };

        for res_json in qdr.results {
            qdr_converted.results.push(QueryResult {
                document_attributes: res_json.document_attributes,
                result: serde_json::from_value(res_json.result)?,
            });
        }

        Ok(qdr_converted)
    }
}

pub struct ListDocumentsRequest {
    hyper_client: HyperClient,
    request: RequestBuilder,
}

impl ListDocumentsRequest {
    pub(crate) fn new(hyper_client: HyperClient, request: RequestBuilder) -> ListDocumentsRequest {
        ListDocumentsRequest { hyper_client, request }
    }

    request_option!(max_item_count, u64, HEADER_MAX_ITEM_COUNT);
    request_bytes_option!(continuation_token, ContinuationToken, HEADER_CONTINUATION);
    request_option!(consistency_level, ConsistencyLevel, HEADER_CONSISTENCY_LEVEL);
    request_bytes_option!(session_token, String, HEADER_SESSION_TOKEN);
    request_bytes_option!(if_none_match, String, header::IF_NONE_MATCH);
    request_bytes_option!(partition_range_id, String, HEADER_DOCUMENTDB_PARTITIONRANGEID);

    pub fn incremental_feed(mut self) -> Self {
        self.request.header(HEADER_A_IM, HeaderValue::from_static("Incremental feed"));
        self
    }

    pub fn execute<T: DeserializeOwned>(mut self) -> impl Future<Item = ListDocumentsResponse<T>, Error = AzureError> {
        future::result(self.request.body(hyper::Body::empty()))
            .from_err()
            .and_then(move |r| check_status_extract_headers_and_body(self.hyper_client.request(r), StatusCode::OK))
            .and_then(|(headers, whole_body)| Self::extract_result::<T>(&whole_body, &headers))
    }

    fn extract_result<T>(body: &[u8], headers: &HeaderMap) -> Result<ListDocumentsResponse<T>, AzureError>
    where
        T: DeserializeOwned,
    {
        debug!("headers == {:?}", headers);

        let ado = ListDocumentsResponseAdditionalHeaders {
            // This match just tries to extract the info and convert it
            // into the correct type. It is complicated because headers
            // can be missing and also because headers.get<T> will return
            // a T reference (&T) so we need to cast it into the
            // correct type and clone it (in this case into a &str that will
            // become a String using to_owned())
            continuation_token: derive_continuation_token(headers),
            // Here we assume the Charge header to always be present.
            // If problems arise we
            // will change the field to be Option(al).
            charge: derive_request_charge(headers),
            etag: headers.get(header::ETAG).and_then(|v| v.to_str().ok()).map(|s| s.to_owned()),
        };
        debug!("ado == {:?}", ado);

        // we will proceed in three steps:
        // 1- Deserialize the result as DocumentAttributes. The extra field will be ignored.
        // 2- Deserialize the result a type T. The extra fields will be ignored.
        // 3- Zip 1 and 2 in the resulting structure.
        // There is a lot of data movement here, let's hope the compiler is smarter than me :)
        let document_attributes = serde_json::from_slice::<ListDocumentsResponseAttributes>(body)?;
        let entries = serde_json::from_slice::<ListDocumentsResponseEntities<T>>(body)?;

        let documents = document_attributes
            .documents
            .into_iter()
            .zip(entries.entities.into_iter())
            .map(|(da, e)| Document {
                document_attributes: da,
                entity: e,
            })
            .collect();

        Ok(ListDocumentsResponse {
            rid: document_attributes.rid,
            documents,
            additional_headers: ado,
        })
    }
}

pub struct ReplaceDocumentRequest<T> {
    hyper_client: HyperClient,
    request: RequestBuilder,
    payload: Result<String, serde_json::Error>,
    _t: PhantomData<T>,
}

impl<T> DocumentRequestExt for ReplaceDocumentRequest<T> {
    fn request(&mut self) -> &mut RequestBuilder {
        &mut self.request
    }
}

impl<T: DeserializeOwned> ReplaceDocumentRequest<T> {
    pub(crate) fn new(
        hyper_client: HyperClient,
        request: RequestBuilder,
        payload: Result<String, serde_json::Error>,
    ) -> ReplaceDocumentRequest<T> {
        ReplaceDocumentRequest {
            hyper_client,
            request,
            payload,
            _t: PhantomData,
        }
    }

    request_bytes_option!(if_match, String, header::IF_MATCH);
    request_option!(indexing_directive, IndexingDirective, HEADER_INDEXING_DIRECTIVE);

    pub fn execute(self) -> impl Future<Item = ReplaceDocumentResponse<T>, Error = AzureError> {
        trace!("get_document called(request == {:?}", self.request);
        let hc = self.hyper_client;
        let mut req = self.request;
        future::result(self.payload)
            .from_err()
            .and_then(move |payload| Ok(req.body(payload.into())?))
            .and_then(move |r| check_status_extract_headers_and_body(hc.request(r), StatusCode::OK))
            .and_then(move |(headers, body)| Self::extract_result(&headers, &body))
    }

    fn extract_result<R: DeserializeOwned>(headers: &HeaderMap, body: &[u8]) -> Result<ReplaceDocumentResponse<R>, AzureError> {
        let additional_headers = DocumentAdditionalHeaders::derive_from(headers);
        let document = Document::from_json(body)?;
        Ok(ReplaceDocumentResponse {
            document,
            additional_headers,
        })
    }
}

pub struct DeleteDocumentRequest {
    hyper_client: HyperClient,
    request: RequestBuilder,
}

impl DocumentRequestExt for DeleteDocumentRequest {
    fn request(&mut self) -> &mut RequestBuilder {
        &mut self.request
    }
}

impl DeleteDocumentRequest {
    pub(crate) fn new(hyper_client: HyperClient, request: RequestBuilder) -> DeleteDocumentRequest {
        DeleteDocumentRequest { hyper_client, request }
    }

    request_bytes_option!(if_match, String, header::IF_MATCH);

    pub fn execute(mut self) -> impl Future<Item = (), Error = AzureError> {
        trace!("get_document called(request == {:?}", self.request);

        future::result(self.request.body(hyper::Body::empty()))
            .from_err()
            .and_then(move |r| check_status_extract_body(self.hyper_client.request(r), StatusCode::NO_CONTENT))
            .and_then(|_| Ok(()))
    }
}

pub trait DocumentRequestExt: Sized {
    fn request(&mut self) -> &mut RequestBuilder;

    fn session_token<S: AsRef<str>>(mut self, token: S) -> Self {
        self.request().header_formatted(HEADER_SESSION_TOKEN, token.as_ref());
        self
    }

    fn partition_key<'a, P: Into<PartitionKey<'a>>>(mut self, key: P) -> Self {
        // todo: move unwrap into PartitionKey impl itself as we control the impl and it surely won't error out
        if let Some(ser_key) = key.into().to_json().unwrap() {
            self.request().header_formatted(HEADER_DOCUMENTDB_PARTITIONKEY, ser_key);
        }
        self
    }
}

fn derive_continuation_token(headers: &HeaderMap) -> Option<String> {
    headers.get(HEADER_CONTINUATION).and_then(|v| v.to_str().ok()).map(|v| v.to_owned())
}

fn derive_request_charge(headers: &HeaderMap) -> f64 {
    headers.get(HEADER_REQUEST_CHARGE).unwrap().to_str().unwrap().parse().unwrap()
}
<|endoftext|>