<commit_before>use serde_json::value::Value as Json;

use crate::context::Context;
use crate::error::RenderError;
use crate::helpers::HelperDef;
use crate::json::value::ScopedJson;
use crate::registry::Registry;
use crate::render::{Helper, RenderContext};

#[derive(Clone, Copy)]
pub struct LookupHelper;

impl HelperDef for LookupHelper {
    fn call_inner<'reg: 'rc, 'rc>(
        &self,
        h: &Helper<'rc>,
        r: &'reg Registry<'reg>,
        _: &'rc Context,
        _: &mut RenderContext<'reg, 'rc>,
    ) -> Result<ScopedJson<'rc>, RenderError> {
        let collection_value = h
            .param(0)
            .ok_or_else(|| RenderError::new("Param not found for helper \"lookup\""))?;
        let index = h
            .param(1)
            .ok_or_else(|| RenderError::new("Insufficient params for helper \"lookup\""))?;

        let value = match *collection_value.value() {
            Json::Array(ref v) => index
                .value()
                .as_u64()
                .and_then(|u| v.get(u as usize))
                .unwrap_or(&Json::Null),
            Json::Object(ref m) => index
                .value()
                .as_str()
                .and_then(|k| m.get(k))
                .unwrap_or(&Json::Null),
            _ => &Json::Null,
        };
        if r.strict_mode() && value.is_null() {
            Err(RenderError::strict_error(None))
        } else {
            Ok(value.clone().into())
        }
    }
}

pub static LOOKUP_HELPER: LookupHelper = LookupHelper;

#[cfg(test)]
mod test {
    use crate::registry::Registry;

    #[test]
    fn test_lookup() {
        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string("t0", "{{#each v1}}{{lookup ../v2 @index}}{{/each}}")
            .is_ok());
        assert!(handlebars
            .register_template_string("t1", "{{#each v1}}{{lookup ../v2 1}}{{/each}}")
            .is_ok());
        assert!(handlebars
            .register_template_string("t2", "{{lookup kk \"a\"}}")
            .is_ok());

        let m = json!({"v1": [1,2,3], "v2": [9,8,7]});

        let m2 = json!({
            "kk": {"a": "world"}
        });

        let r0 = handlebars.render("t0", &m);
        assert_eq!(r0.ok().unwrap(), "987".to_string());

        let r1 = handlebars.render("t1", &m);
        assert_eq!(r1.ok().unwrap(), "888".to_string());

        let r2 = handlebars.render("t2", &m2);
        assert_eq!(r2.ok().unwrap(), "world".to_string());

        assert!(handlebars.render_template("{{lookup}}", &m).is_err());
        assert!(handlebars.render_template("{{lookup v1}}", &m).is_err());
        assert_eq!(
            handlebars.render_template("{{lookup null 1}}", &m).unwrap(),
            ""
        );
        assert_eq!(
            handlebars.render_template("{{lookup v1 3}}", &m).unwrap(),
            ""
        );
    }

    #[test]
    fn test_strict_lookup() {
        let mut hbs = Registry::new();

        assert_eq!(
            hbs.render_template("{{lookup kk 1}}", &json!({"kk": []}))
                .unwrap(),
            ""
        );

        hbs.set_strict_mode(true);

        assert!(hbs
            .render_template("{{lookup kk 1}}", &json!({"kk": []}))
            .is_err());
    }
}
<commit_msg>fix: distiguish null and missing value in lookup<commit_after>use serde_json::value::Value as Json;

use crate::context::Context;
use crate::error::RenderError;
use crate::helpers::HelperDef;
use crate::json::value::ScopedJson;
use crate::registry::Registry;
use crate::render::{Helper, RenderContext};

#[derive(Clone, Copy)]
pub struct LookupHelper;

impl HelperDef for LookupHelper {
    fn call_inner<'reg: 'rc, 'rc>(
        &self,
        h: &Helper<'rc>,
        r: &'reg Registry<'reg>,
        _: &'rc Context,
        _: &mut RenderContext<'reg, 'rc>,
    ) -> Result<ScopedJson<'rc>, RenderError> {
        let collection_value = h
            .param(0)
            .ok_or_else(|| RenderError::new("Param not found for helper \"lookup\""))?;
        let index = h
            .param(1)
            .ok_or_else(|| RenderError::new("Insufficient params for helper \"lookup\""))?;

        let value = match *collection_value.value() {
            Json::Array(ref v) => index.value().as_u64().and_then(|u| v.get(u as usize)),
            Json::Object(ref m) => index.value().as_str().and_then(|k| m.get(k)),
            _ => None,
        };
        if r.strict_mode() && value.is_none() {
            Err(RenderError::strict_error(None))
        } else {
            Ok(value.unwrap_or(&Json::Null).clone().into())
        }
    }
}

pub static LOOKUP_HELPER: LookupHelper = LookupHelper;

#[cfg(test)]
mod test {
    use crate::registry::Registry;

    #[test]
    fn test_lookup() {
        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string("t0", "{{#each v1}}{{lookup ../v2 @index}}{{/each}}")
            .is_ok());
        assert!(handlebars
            .register_template_string("t1", "{{#each v1}}{{lookup ../v2 1}}{{/each}}")
            .is_ok());
        assert!(handlebars
            .register_template_string("t2", "{{lookup kk \"a\"}}")
            .is_ok());

        let m = json!({"v1": [1,2,3], "v2": [9,8,7]});

        let m2 = json!({
            "kk": {"a": "world"}
        });

        let r0 = handlebars.render("t0", &m);
        assert_eq!(r0.ok().unwrap(), "987".to_string());

        let r1 = handlebars.render("t1", &m);
        assert_eq!(r1.ok().unwrap(), "888".to_string());

        let r2 = handlebars.render("t2", &m2);
        assert_eq!(r2.ok().unwrap(), "world".to_string());

        assert!(handlebars.render_template("{{lookup}}", &m).is_err());
        assert!(handlebars.render_template("{{lookup v1}}", &m).is_err());
        assert_eq!(
            handlebars.render_template("{{lookup null 1}}", &m).unwrap(),
            ""
        );
        assert_eq!(
            handlebars.render_template("{{lookup v1 3}}", &m).unwrap(),
            ""
        );
    }

    #[test]
    fn test_strict_lookup() {
        let mut hbs = Registry::new();

        assert_eq!(
            hbs.render_template("{{lookup kk 1}}", &json!({"kk": []}))
                .unwrap(),
            ""
        );
        assert!(hbs
            .render_template("{{lookup kk 0}}", &json!({ "kk": [null] }))
            .is_ok());

        hbs.set_strict_mode(true);

        assert!(hbs
            .render_template("{{lookup kk 1}}", &json!({"kk": []}))
            .is_err());
        assert!(hbs
            .render_template("{{lookup kk 0}}", &json!({ "kk": [null] }))
            .is_ok());
    }
}
<|endoftext|>


<commit_before>use sozu::channel::Channel;
use sozu::messages::{Order, Instance};
use sozu_command::data::{AnswerData,ConfigCommand,ConfigMessage,ConfigMessageAnswer,ConfigMessageStatus,RunState};

use std::collections::HashSet;
use rand::{thread_rng, Rng};

fn generate_id() -> String {
  let s: String = thread_rng().gen_ascii_chars().take(6).collect();
  format!("ID-{}", s)
}

pub fn save_state(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, path: &str) {
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::SaveState(path.to_string()),
    None,
  ));

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          // do nothing here
          // for other messages, we would loop over read_message
          // until an error or ok message was sent
        },
        ConfigMessageStatus::Error => {
          println!("could not save proxy state: {}", message.message);
        },
        ConfigMessageStatus::Ok => {
          println!("Proxy state saved to {}", path);
        }
      }
    }
  }
}

pub fn load_state(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, path: &str) {
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::LoadState(path.to_string()),
    None,
  ));

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          // do nothing here
          // for other messages, we would loop over read_message
          // until an error or ok message was sent
        },
        ConfigMessageStatus::Error => {
          println!("could not save proxy state: {}", message.message);
        },
        ConfigMessageStatus::Ok => {
          println!("Proxy state saved to {}", path);
        }
      }
    }
  }
}

pub fn dump_state(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>) {
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::DumpState,
    None,
  ));

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          // do nothing here
          // for other messages, we would loop over read_message
          // until an error or ok message was sent
        },
        ConfigMessageStatus::Error => {
          println!("could not dump proxy state: {}", message.message);
        },
        ConfigMessageStatus::Ok => {
          println!("Proxy state:\n{}", message.message);
        }
      }
    }
  }
}

pub fn soft_stop(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>) {
  println!("shutting down proxy");
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::ProxyConfiguration(Order::SoftStop),
    //FIXME: we should be able to soft stop one specific worker
    None,
  ));

  loop {
    match channel.read_message() {
      None          => println!("the proxy didn't answer"),
      Some(message) => {
        if &id != &message.id {
          println!("received message with invalid id: {:?}", message);
          return;
        }
        match message.status {
          ConfigMessageStatus::Processing => {
            println!("Proxy is processing: {}", message.message);
          },
          ConfigMessageStatus::Error => {
            println!("could not stop the proxy: {}", message.message);
          },
          ConfigMessageStatus::Ok => {
            println!("Proxy shut down: {}", message.message);
            break;
          }
        }
      }
    }
  }
}

pub fn hard_stop(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>) {
  println!("shutting down proxy");
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::ProxyConfiguration(Order::HardStop),
    //FIXME: we should be able to soft stop one specific worker
    None,
  ));

  loop {
    match channel.read_message() {
      None          => println!("the proxy didn't answer"),
      Some(message) => {
        match message.status {
          ConfigMessageStatus::Processing => {
            println!("Proxy is processing: {}", message.message);
          },
          ConfigMessageStatus::Error => {
            println!("could not stop the proxy: {}", message.message);
          },
          ConfigMessageStatus::Ok => {
            if &id == &message.id {
              println!("Proxy shut down: {}", message.message);
              break;
            }
          }
        }
      }
    }
  }
}

pub fn upgrade(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>) {
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::ListWorkers,
    None,
  ));

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          println!("should have obtained an answer immediately");
          return;
        },
        ConfigMessageStatus::Error => {
          println!("could not get the worker list: {}", message.message);
          return
        },
        ConfigMessageStatus::Ok => {
          println!("Worker list:\n{:?}", message.data);
          if let Some(AnswerData::Workers(ref workers)) = message.data {
            let mut launching: HashSet<String> = HashSet::new();
            let mut stopping:  HashSet<String> = HashSet::new();

            for ref worker in workers.iter().filter(|worker| worker.run_state == RunState::Running) {
              let id = generate_id();
              let msg = ConfigMessage::new(
                id.clone(),
                ConfigCommand::LaunchWorker("BLAH".to_string()),
                None,
              );
              println!("sending message: {:?}", msg);
              channel.write_message(&msg);
              launching.insert(id);
            }

            for ref worker in workers.iter().filter(|worker| worker.run_state == RunState::Running) {
              let id = generate_id();
              let msg = ConfigMessage::new(
                id.clone(),
                ConfigCommand::ProxyConfiguration(Order::SoftStop),
                Some(worker.id),
              );
              println!("sending message: {:?}", msg);
              channel.write_message(&msg);
              stopping.insert(id);
            }

            loop {
              println!("launching: {:?}\nstopping: {:?}", launching, stopping);
              if launching.is_empty() && stopping.is_empty() {
                break;
              }
              match channel.read_message() {
                None          => println!("the proxy didn't answer"),
                Some(message) => {
                  println!("received message: {:?}", message);
                  match message.status {
                    ConfigMessageStatus::Processing => {
                    },
                    ConfigMessageStatus::Error => {
                      println!("error for message[{}]: {}", message.id, message.message);
                      if launching.contains(&message.id) {
                        launching.remove(&message.id);
                        println!("launch message with ID {} done", message.id);
                      }
                      if stopping.contains(&message.id) {
                        stopping.remove(&message.id);
                        println!("stop message with ID {} done", message.id);
                      }
                    },
                    ConfigMessageStatus::Ok => {
                      if launching.contains(&message.id) {
                        launching.remove(&message.id);
                        println!("launch message with ID {} done", message.id);
                      }
                      if stopping.contains(&message.id) {
                        stopping.remove(&message.id);
                        println!("stop message with ID {} done", message.id);
                      }
                    }
                  }
                }
              }
            }

            println!("worker upgrade done");
            let id = generate_id();
            channel.write_message(&ConfigMessage::new(
              id.clone(),
              ConfigCommand::UpgradeMaster,
              None,
            ));

            loop {
              match channel.read_message() {
                None          => println!("the proxy didn't answer"),
                Some(message) => {
                  if &id != &message.id {
                    println!("received message with invalid id: {:?}", message);
                    return;
                  }
                  match message.status {
                    ConfigMessageStatus::Processing => {
                      println!("master is processing: {}", message.message);
                    },
                    ConfigMessageStatus::Error => {
                      println!("could not upgrade the master: {}", message.message);
                      return;
                    },
                    ConfigMessageStatus::Ok => {
                      println!("successfully upgraded the master: {}", message.message);
                      return;
                    }
                  }
                }
              }
            }
          }
        }
      }
    }
  }
}

pub fn status(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>) {
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::ListWorkers,
    None,
  ));

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          println!("should have obtained an answer immediately");
          return;
        },
        ConfigMessageStatus::Error => {
          println!("could not get the worker list: {}", message.message);
          return
        },
        ConfigMessageStatus::Ok => {
          println!("Worker list:\n{:?}", message.data);
          if let Some(AnswerData::Workers(ref workers)) = message.data {
            let mut expecting: HashSet<String> = HashSet::new();

            for ref worker in workers.iter().filter(|worker| worker.run_state == RunState::Running) {
              let id = generate_id();
              let msg = ConfigMessage::new(
                id.clone(),
                ConfigCommand::ProxyConfiguration(Order::Status),
                Some(worker.id),
              );
              println!("sending message: {:?}", msg);
              channel.write_message(&msg);
              expecting.insert(id);
            }

            loop {
              println!("expecting: {:?}", expecting);
              if expecting.is_empty() {
                break;
              }
              match channel.read_message() {
                None          => println!("the proxy didn't answer"),
                Some(message) => {
                  println!("received message: {:?}", message);
                  match message.status {
                    ConfigMessageStatus::Processing => {
                    },
                    ConfigMessageStatus::Error => {
                      println!("error for message[{}]: {}", message.id, message.message);
                      if expecting.contains(&message.id) {
                        expecting.remove(&message.id);
                        println!("status message with ID {} done", message.id);
                      }
                    },
                    ConfigMessageStatus::Ok => {
                      if expecting.contains(&message.id) {
                        expecting.remove(&message.id);
                        println!("status message with ID {} done", message.id);
                      }
                    }
                  }
                }
              }
            }

            println!("worker upgrade done");
          }
        }
      }
    }
  }
}

pub fn remove_backend(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, app_id: &str, ip: &str, port: u16) {
    let instance = Instance {
      app_id: String::from(app_id),
      ip_address: String::from(ip),
      port: port
    };

    backend_command(channel, Order::RemoveInstance(instance.clone()), instance, "remove");
}

pub fn add_backend(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, app_id: &str, ip: &str, port: u16) {
  let instance = Instance {
      app_id: String::from(app_id),
      ip_address: String::from(ip),
      port: port
    };

  backend_command(channel, Order::AddInstance(instance.clone()), instance, "add");
}

fn backend_command(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, order: Order, instance: Instance, action: &str) {
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::ProxyConfiguration(order),
    None,
  ));

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          // do nothing here
          // for other messages, we would loop over read_message
          // until an error or ok message was sent
        },
        ConfigMessageStatus::Error => {
          println!("could not {} backend : {}", action, message.message);
        },
        ConfigMessageStatus::Ok => {
          println!("backend {}:{} {} for app : {} ", instance.ip_address, instance.port, action, instance.app_id);
        }
      }
    }
  }
}
<commit_msg>refactor thing<commit_after>use sozu::channel::Channel;
use sozu::messages::{Order, Instance};
use sozu_command::data::{AnswerData,ConfigCommand,ConfigMessage,ConfigMessageAnswer,ConfigMessageStatus,RunState};

use std::collections::HashSet;
use rand::{thread_rng, Rng};

fn generate_id() -> String {
  let s: String = thread_rng().gen_ascii_chars().take(6).collect();
  format!("ID-{}", s)
}

pub fn save_state(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, path: &str) {
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::SaveState(path.to_string()),
    None,
  ));

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          // do nothing here
          // for other messages, we would loop over read_message
          // until an error or ok message was sent
        },
        ConfigMessageStatus::Error => {
          println!("could not save proxy state: {}", message.message);
        },
        ConfigMessageStatus::Ok => {
          println!("Proxy state saved to {}", path);
        }
      }
    }
  }
}

pub fn load_state(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, path: &str) {
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::LoadState(path.to_string()),
    None,
  ));

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          // do nothing here
          // for other messages, we would loop over read_message
          // until an error or ok message was sent
        },
        ConfigMessageStatus::Error => {
          println!("could not save proxy state: {}", message.message);
        },
        ConfigMessageStatus::Ok => {
          println!("Proxy state saved to {}", path);
        }
      }
    }
  }
}

pub fn dump_state(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>) {
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::DumpState,
    None,
  ));

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          // do nothing here
          // for other messages, we would loop over read_message
          // until an error or ok message was sent
        },
        ConfigMessageStatus::Error => {
          println!("could not dump proxy state: {}", message.message);
        },
        ConfigMessageStatus::Ok => {
          println!("Proxy state:\n{}", message.message);
        }
      }
    }
  }
}

pub fn soft_stop(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>) {
  println!("shutting down proxy");
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::ProxyConfiguration(Order::SoftStop),
    //FIXME: we should be able to soft stop one specific worker
    None,
  ));

  loop {
    match channel.read_message() {
      None          => println!("the proxy didn't answer"),
      Some(message) => {
        if &id != &message.id {
          println!("received message with invalid id: {:?}", message);
          return;
        }
        match message.status {
          ConfigMessageStatus::Processing => {
            println!("Proxy is processing: {}", message.message);
          },
          ConfigMessageStatus::Error => {
            println!("could not stop the proxy: {}", message.message);
          },
          ConfigMessageStatus::Ok => {
            println!("Proxy shut down: {}", message.message);
            break;
          }
        }
      }
    }
  }
}

pub fn hard_stop(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>) {
  println!("shutting down proxy");
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::ProxyConfiguration(Order::HardStop),
    //FIXME: we should be able to soft stop one specific worker
    None,
  ));

  loop {
    match channel.read_message() {
      None          => println!("the proxy didn't answer"),
      Some(message) => {
        match message.status {
          ConfigMessageStatus::Processing => {
            println!("Proxy is processing: {}", message.message);
          },
          ConfigMessageStatus::Error => {
            println!("could not stop the proxy: {}", message.message);
          },
          ConfigMessageStatus::Ok => {
            if &id == &message.id {
              println!("Proxy shut down: {}", message.message);
              break;
            }
          }
        }
      }
    }
  }
}

pub fn upgrade(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>) {
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::ListWorkers,
    None,
  ));

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          println!("should have obtained an answer immediately");
          return;
        },
        ConfigMessageStatus::Error => {
          println!("could not get the worker list: {}", message.message);
          return
        },
        ConfigMessageStatus::Ok => {
          println!("Worker list:\n{:?}", message.data);
          if let Some(AnswerData::Workers(ref workers)) = message.data {
            let mut launching: HashSet<String> = HashSet::new();
            let mut stopping:  HashSet<String> = HashSet::new();

            for ref worker in workers.iter().filter(|worker| worker.run_state == RunState::Running) {
              let id = generate_id();
              let msg = ConfigMessage::new(
                id.clone(),
                ConfigCommand::LaunchWorker("BLAH".to_string()),
                None,
              );
              println!("sending message: {:?}", msg);
              channel.write_message(&msg);
              launching.insert(id);
            }

            for ref worker in workers.iter().filter(|worker| worker.run_state == RunState::Running) {
              let id = generate_id();
              let msg = ConfigMessage::new(
                id.clone(),
                ConfigCommand::ProxyConfiguration(Order::SoftStop),
                Some(worker.id),
              );
              println!("sending message: {:?}", msg);
              channel.write_message(&msg);
              stopping.insert(id);
            }

            loop {
              println!("launching: {:?}\nstopping: {:?}", launching, stopping);
              if launching.is_empty() && stopping.is_empty() {
                break;
              }
              match channel.read_message() {
                None          => println!("the proxy didn't answer"),
                Some(message) => {
                  println!("received message: {:?}", message);
                  match message.status {
                    ConfigMessageStatus::Processing => {
                    },
                    ConfigMessageStatus::Error => {
                      println!("error for message[{}]: {}", message.id, message.message);
                      if launching.contains(&message.id) {
                        launching.remove(&message.id);
                        println!("launch message with ID {} done", message.id);
                      }
                      if stopping.contains(&message.id) {
                        stopping.remove(&message.id);
                        println!("stop message with ID {} done", message.id);
                      }
                    },
                    ConfigMessageStatus::Ok => {
                      if launching.contains(&message.id) {
                        launching.remove(&message.id);
                        println!("launch message with ID {} done", message.id);
                      }
                      if stopping.contains(&message.id) {
                        stopping.remove(&message.id);
                        println!("stop message with ID {} done", message.id);
                      }
                    }
                  }
                }
              }
            }

            println!("worker upgrade done");
            let id = generate_id();
            channel.write_message(&ConfigMessage::new(
              id.clone(),
              ConfigCommand::UpgradeMaster,
              None,
            ));

            loop {
              match channel.read_message() {
                None          => println!("the proxy didn't answer"),
                Some(message) => {
                  if &id != &message.id {
                    println!("received message with invalid id: {:?}", message);
                    return;
                  }
                  match message.status {
                    ConfigMessageStatus::Processing => {
                      println!("master is processing: {}", message.message);
                    },
                    ConfigMessageStatus::Error => {
                      println!("could not upgrade the master: {}", message.message);
                      return;
                    },
                    ConfigMessageStatus::Ok => {
                      println!("successfully upgraded the master: {}", message.message);
                      return;
                    }
                  }
                }
              }
            }
          }
        }
      }
    }
  }
}

pub fn status(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>) {
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::ListWorkers,
    None,
  ));

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          println!("should have obtained an answer immediately");
          return;
        },
        ConfigMessageStatus::Error => {
          println!("could not get the worker list: {}", message.message);
          return
        },
        ConfigMessageStatus::Ok => {
          println!("Worker list:\n{:?}", message.data);
          if let Some(AnswerData::Workers(ref workers)) = message.data {
            let mut expecting: HashSet<String> = HashSet::new();

            for ref worker in workers.iter().filter(|worker| worker.run_state == RunState::Running) {
              let id = generate_id();
              let msg = ConfigMessage::new(
                id.clone(),
                ConfigCommand::ProxyConfiguration(Order::Status),
                Some(worker.id),
              );
              println!("sending message: {:?}", msg);
              channel.write_message(&msg);
              expecting.insert(id);
            }

            loop {
              println!("expecting: {:?}", expecting);
              if expecting.is_empty() {
                break;
              }
              match channel.read_message() {
                None          => println!("the proxy didn't answer"),
                Some(message) => {
                  println!("received message: {:?}", message);
                  match message.status {
                    ConfigMessageStatus::Processing => {
                    },
                    ConfigMessageStatus::Error => {
                      println!("error for message[{}]: {}", message.id, message.message);
                      if expecting.contains(&message.id) {
                        expecting.remove(&message.id);
                        println!("status message with ID {} done", message.id);
                      }
                    },
                    ConfigMessageStatus::Ok => {
                      if expecting.contains(&message.id) {
                        expecting.remove(&message.id);
                        println!("status message with ID {} done", message.id);
                      }
                    }
                  }
                }
              }
            }

            println!("worker upgrade done");
          }
        }
      }
    }
  }
}

pub fn remove_backend(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, app_id: &str, ip: &str, port: u16) {
    let instance = Instance {
      app_id: String::from(app_id),
      ip_address: String::from(ip),
      port: port
    };

    order_command(channel, Order::RemoveInstance(instance));
}

pub fn add_backend(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, app_id: &str, ip: &str, port: u16) {
  let instance = Instance {
      app_id: String::from(app_id),
      ip_address: String::from(ip),
      port: port
    };

  order_command(channel, Order::AddInstance(instance));
}

fn order_command(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, order: Order) {
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::ProxyConfiguration(order.clone()),
    None,
  ));

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          // do nothing here
          // for other messages, we would loop over read_message
          // until an error or ok message was sent
        },
        ConfigMessageStatus::Error => {
          match order {
            Order::AddInstance(_) => println!("could not add backend : {}", message.message),
            Order::RemoveInstance(_) => println!("could not remove backend : {}", message.message),
            _ => {
              // do nothing for now
            }
          }
        },
        ConfigMessageStatus::Ok => {
          match order {
            Order::AddInstance(i) => println!("backend {}:{} added for app : {} ", i.ip_address, i.port, i.app_id),
            Order::RemoveInstance(i) => println!("backend {}:{} removed for app : {} ", i.ip_address, i.port, i.app_id),
            _ => {
              // do nothing for now
            }
          }
        }
      }
    }
  }
}
<|endoftext|>


<commit_before>// Copyright (C) 2018 Sebastian Dröge <sebastian@centricular.com>
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Library General Public
// License as published by the Free Software Foundation; either
// version 2 of the License, or (at your option) any later version.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Library General Public License for more details.
//
// You should have received a copy of the GNU Library General Public
// License along with this library; if not, write to the
// Free Software Foundation, Inc., 51 Franklin Street, Suite 500,
// Boston, MA 02110-1335, USA.

use std::collections::HashMap;
use std::io;
use std::mem;
use std::sync::{atomic, mpsc};
use std::sync::{Arc, Mutex, Weak};
use std::thread;

use futures::future;
use futures::stream::futures_unordered::FuturesUnordered;
use futures::sync::oneshot;
use futures::{Future, Stream};
use tokio::reactor;
use tokio_current_thread;
use tokio_timer::timer;

use glib;
use gst;

lazy_static! {
    static ref CONTEXTS: Mutex<HashMap<String, Weak<IOContextInner>>> = Mutex::new(HashMap::new());
    static ref CONTEXT_CAT: gst::DebugCategory = gst::DebugCategory::new(
        "ts-context",
        gst::DebugColorFlags::empty(),
        Some("Thread-sharing Context"),
    );
}

// Our own simplified implementation of reactor::Background to allow hooking into its internals
const RUNNING: usize = 0;
const SHUTDOWN_NOW: usize = 1;

struct IOContextRunner {
    name: String,
    shutdown: Arc<atomic::AtomicUsize>,
}

impl IOContextRunner {
    fn start(
        name: &str,
        wait: u32,
        reactor: reactor::Reactor,
    ) -> (tokio_current_thread::Handle, IOContextShutdown) {
        let handle = reactor.handle().clone();
        let shutdown = Arc::new(atomic::AtomicUsize::new(RUNNING));
        let shutdown_clone = shutdown.clone();
        let name_clone = name.into();

        let mut runner = IOContextRunner {
            shutdown: shutdown_clone,
            name: name_clone,
        };

        let (sender, receiver) = mpsc::channel();

        let join = thread::spawn(move || {
            runner.run(wait, reactor, sender);
        });

        let shutdown = IOContextShutdown {
            name: name.into(),
            shutdown,
            handle,
            join: Some(join),
        };

        let runtime_handle = receiver.recv().unwrap();

        (runtime_handle, shutdown)
    }

    fn run(
        &mut self,
        wait: u32,
        reactor: reactor::Reactor,
        sender: mpsc::Sender<tokio_current_thread::Handle>,
    ) {
        use std::time;
        let wait = time::Duration::from_millis(wait as u64);

        gst_debug!(CONTEXT_CAT, "Started reactor thread '{}'", self.name);

        let handle = reactor.handle();
        let mut enter = ::tokio_executor::enter().unwrap();
        let timer = timer::Timer::new(reactor);
        let timer_handle = timer.handle();
        let mut current_thread = tokio_current_thread::CurrentThread::new_with_park(timer);

        let _ = sender.send(current_thread.handle());

        let mut now = time::Instant::now();

        ::tokio_timer::with_default(&timer_handle, &mut enter, |mut enter| {
            ::tokio_reactor::with_default(&handle, &mut enter, |enter| loop {
                if self.shutdown.load(atomic::Ordering::SeqCst) > RUNNING {
                    break;
                }

                gst_trace!(CONTEXT_CAT, "Turning current thread '{}'", self.name);
                while current_thread
                    .enter(enter)
                    .turn(Some(time::Duration::from_millis(0)))
                    .unwrap()
                    .has_polled()
                {}
                gst_trace!(CONTEXT_CAT, "Turned current thread '{}'", self.name);

                let elapsed = now.elapsed();
                if elapsed < wait {
                    gst_trace!(
                        CONTEXT_CAT,
                        "Waiting for {:?} before polling again",
                        wait - elapsed
                    );
                    thread::sleep(wait - elapsed);
                }
                now += wait;
            })
        });
    }
}

impl Drop for IOContextRunner {
    fn drop(&mut self) {
        gst_debug!(CONTEXT_CAT, "Shut down reactor thread '{}'", self.name);
    }
}

struct IOContextShutdown {
    name: String,
    shutdown: Arc<atomic::AtomicUsize>,
    handle: reactor::Handle,
    join: Option<thread::JoinHandle<()>>,
}

impl Drop for IOContextShutdown {
    fn drop(&mut self) {
        use tokio_executor::park::Unpark;

        gst_debug!(CONTEXT_CAT, "Shutting down reactor thread '{}'", self.name);
        self.shutdown.store(SHUTDOWN_NOW, atomic::Ordering::SeqCst);
        gst_trace!(CONTEXT_CAT, "Waiting for reactor '{}' shutdown", self.name);
        // After being unparked, the next turn() is guaranteed to finish immediately,
        // as such there is no race condition between checking for shutdown and setting
        // shutdown.
        self.handle.unpark();
        let _ = self.join.take().unwrap().join();
    }
}

#[derive(Clone)]
pub struct IOContext(Arc<IOContextInner>);

impl glib::subclass::boxed::BoxedType for IOContext {
    const NAME: &'static str = "TsIOContext";

    glib_boxed_type!();
}

glib_boxed_derive_traits!(IOContext);

struct IOContextInner {
    name: String,
    runtime_handle: Mutex<tokio_current_thread::Handle>,
    reactor_handle: reactor::Handle,
    // Only used for dropping
    _shutdown: IOContextShutdown,
    pending_futures: Mutex<(
        u64,
        HashMap<u64, FuturesUnordered<Box<dyn Future<Item = (), Error = ()> + Send + 'static>>>,
    )>,
}

impl Drop for IOContextInner {
    fn drop(&mut self) {
        let mut contexts = CONTEXTS.lock().unwrap();
        gst_debug!(CONTEXT_CAT, "Finalizing context '{}'", self.name);
        contexts.remove(&self.name);
    }
}

impl IOContext {
    pub fn new(name: &str, wait: u32) -> Result<Self, io::Error> {
        let mut contexts = CONTEXTS.lock().unwrap();
        if let Some(context) = contexts.get(name) {
            if let Some(context) = context.upgrade() {
                gst_debug!(CONTEXT_CAT, "Reusing existing context '{}'", name);
                return Ok(IOContext(context));
            }
        }

        let reactor = reactor::Reactor::new()?;
        let reactor_handle = reactor.handle().clone();

        let (runtime_handle, shutdown) = IOContextRunner::start(name, wait, reactor);

        let context = Arc::new(IOContextInner {
            name: name.into(),
            runtime_handle: Mutex::new(runtime_handle),
            reactor_handle,
            _shutdown: shutdown,
            pending_futures: Mutex::new((0, HashMap::new())),
        });
        contexts.insert(name.into(), Arc::downgrade(&context));

        gst_debug!(CONTEXT_CAT, "Created new context '{}'", name);
        Ok(IOContext(context))
    }

    pub fn spawn<F>(&self, future: F)
    where
        F: Future<Item = (), Error = ()> + Send + 'static,
    {
        self.0.runtime_handle.lock().unwrap().spawn(future).unwrap();
    }

    pub fn reactor_handle(&self) -> &reactor::Handle {
        &self.0.reactor_handle
    }

    pub fn acquire_pending_future_id(&self) -> PendingFutureId {
        let mut pending_futures = self.0.pending_futures.lock().unwrap();
        let id = pending_futures.0;
        pending_futures.0 += 1;
        pending_futures.1.insert(id, FuturesUnordered::new());

        PendingFutureId(id)
    }

    pub fn release_pending_future_id(&self, id: PendingFutureId) {
        let mut pending_futures = self.0.pending_futures.lock().unwrap();
        if let Some(fs) = pending_futures.1.remove(&id.0) {
            self.spawn(fs.for_each(|_| Ok(())));
        }
    }

    pub fn add_pending_future<F>(&self, id: PendingFutureId, future: F)
    where
        F: Future<Item = (), Error = ()> + Send + 'static,
    {
        let mut pending_futures = self.0.pending_futures.lock().unwrap();
        let fs = pending_futures.1.get_mut(&id.0).unwrap();
        fs.push(Box::new(future))
    }

    pub fn drain_pending_futures<E: Send + 'static>(
        &self,
        id: PendingFutureId,
    ) -> (Option<oneshot::Sender<()>>, PendingFuturesFuture<E>) {
        let mut pending_futures = self.0.pending_futures.lock().unwrap();
        let fs = pending_futures.1.get_mut(&id.0).unwrap();

        let pending_futures = mem::replace(fs, FuturesUnordered::new());

        if !pending_futures.is_empty() {
            gst_log!(
                CONTEXT_CAT,
                "Scheduling {} pending futures for context '{}' with pending future id {:?}",
                pending_futures.len(),
                self.0.name,
                id,
            );

            let (sender, receiver) = oneshot::channel();

            let future = pending_futures
                .for_each(|_| Ok(()))
                .select(receiver.then(|_| Ok(())))
                .then(|_| Ok(()));

            (Some(sender), future::Either::A(Box::new(future)))
        } else {
            (None, future::Either::B(future::ok(())))
        }
    }
}

pub type PendingFuturesFuture<E> = future::Either<
    Box<dyn Future<Item = (), Error = E> + Send + 'static>,
    future::FutureResult<(), E>,
>;

#[derive(Clone, Copy, Eq, PartialEq, Hash, Debug)]
pub struct PendingFutureId(u64);

impl glib::subclass::boxed::BoxedType for PendingFutureId {
    const NAME: &'static str = "TsPendingFutureId";

    glib_boxed_type!();
}

glib_boxed_derive_traits!(PendingFutureId);
<commit_msg>threadshare/iocontext: Add custom interval timer implementation<commit_after>// Copyright (C) 2018 Sebastian Dröge <sebastian@centricular.com>
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Library General Public
// License as published by the Free Software Foundation; either
// version 2 of the License, or (at your option) any later version.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Library General Public License for more details.
//
// You should have received a copy of the GNU Library General Public
// License along with this library; if not, write to the
// Free Software Foundation, Inc., 51 Franklin Street, Suite 500,
// Boston, MA 02110-1335, USA.

use std::cmp;
use std::collections::{BinaryHeap, HashMap};
use std::io;
use std::mem;
use std::sync::{atomic, mpsc};
use std::sync::{Arc, Mutex, Weak};
use std::thread;
use std::time;

use futures::future;
use futures::stream::futures_unordered::FuturesUnordered;
use futures::sync::mpsc as futures_mpsc;
use futures::sync::oneshot;
use futures::{Future, Stream};
use tokio::reactor;
use tokio_current_thread;
use tokio_timer::timer;

use glib;
use gst;

lazy_static! {
    static ref CONTEXTS: Mutex<HashMap<String, Weak<IOContextInner>>> = Mutex::new(HashMap::new());
    static ref CONTEXT_CAT: gst::DebugCategory = gst::DebugCategory::new(
        "ts-context",
        gst::DebugColorFlags::empty(),
        Some("Thread-sharing Context"),
    );
}

// Our own simplified implementation of reactor::Background to allow hooking into its internals
const RUNNING: usize = 0;
const SHUTDOWN_NOW: usize = 1;

struct IOContextRunner {
    name: String,
    shutdown: Arc<atomic::AtomicUsize>,
}

impl IOContextRunner {
    fn start(
        name: &str,
        wait: u32,
        reactor: reactor::Reactor,
        timers: Arc<Mutex<BinaryHeap<TimerEntry>>>,
    ) -> (tokio_current_thread::Handle, IOContextShutdown) {
        let handle = reactor.handle().clone();
        let shutdown = Arc::new(atomic::AtomicUsize::new(RUNNING));
        let shutdown_clone = shutdown.clone();
        let name_clone = name.into();

        let mut runner = IOContextRunner {
            shutdown: shutdown_clone,
            name: name_clone,
        };

        let (sender, receiver) = mpsc::channel();

        let join = thread::spawn(move || {
            runner.run(wait, reactor, sender, timers);
        });

        let shutdown = IOContextShutdown {
            name: name.into(),
            shutdown,
            handle,
            join: Some(join),
        };

        let runtime_handle = receiver.recv().unwrap();

        (runtime_handle, shutdown)
    }

    fn run(
        &mut self,
        wait: u32,
        reactor: reactor::Reactor,
        sender: mpsc::Sender<tokio_current_thread::Handle>,
        timers: Arc<Mutex<BinaryHeap<TimerEntry>>>,
    ) {
        let wait = time::Duration::from_millis(wait as u64);

        gst_debug!(CONTEXT_CAT, "Started reactor thread '{}'", self.name);

        let handle = reactor.handle();
        let mut enter = ::tokio_executor::enter().unwrap();
        let timer = timer::Timer::new(reactor);
        let timer_handle = timer.handle();
        let mut current_thread = tokio_current_thread::CurrentThread::new_with_park(timer);

        let _ = sender.send(current_thread.handle());

        let mut now = time::Instant::now();

        ::tokio_timer::with_default(&timer_handle, &mut enter, |mut enter| {
            ::tokio_reactor::with_default(&handle, &mut enter, |enter| loop {
                if self.shutdown.load(atomic::Ordering::SeqCst) > RUNNING {
                    break;
                }

                gst_trace!(CONTEXT_CAT, "Elapsed {:?} since last loop", now.elapsed());

                // Handle timers
                {
                    // Trigger all timers that would be expired before the middle of the loop wait
                    // time
                    let timer_threshold = now + wait / 2;
                    let mut timers = timers.lock().unwrap();
                    while timers
                        .peek()
                        .and_then(|entry| {
                            if entry.time < timer_threshold {
                                Some(())
                            } else {
                                None
                            }
                        })
                        .is_some()
                    {
                        let TimerEntry {
                            time,
                            interval,
                            sender,
                            ..
                        } = timers.pop().unwrap();

                        if sender.is_closed() {
                            continue;
                        }

                        let _ = sender.unbounded_send(());
                        if let Some(interval) = interval {
                            timers.push(TimerEntry {
                                time: time + interval,
                                id: TIMER_ENTRY_ID.fetch_add(1, atomic::Ordering::Relaxed),
                                interval: Some(interval),
                                sender,
                            });
                        }
                    }
                }

                gst_trace!(CONTEXT_CAT, "Turning current thread '{}'", self.name);
                while current_thread
                    .enter(enter)
                    .turn(Some(time::Duration::from_millis(0)))
                    .unwrap()
                    .has_polled()
                {}
                gst_trace!(CONTEXT_CAT, "Turned current thread '{}'", self.name);

                let elapsed = now.elapsed();
                gst_trace!(CONTEXT_CAT, "Elapsed {:?} after handling futures", elapsed);

                if elapsed < wait {
                    gst_trace!(
                        CONTEXT_CAT,
                        "Waiting for {:?} before polling again",
                        wait - elapsed
                    );
                    thread::sleep(wait - elapsed);
                    gst_trace!(CONTEXT_CAT, "Slept for {:?}", now.elapsed());
                }

                if wait > time::Duration::from_millis(0) {
                    now += wait;
                } else {
                    now = time::Instant::now();
                }
            })
        });
    }
}

impl Drop for IOContextRunner {
    fn drop(&mut self) {
        gst_debug!(CONTEXT_CAT, "Shut down reactor thread '{}'", self.name);
    }
}

struct IOContextShutdown {
    name: String,
    shutdown: Arc<atomic::AtomicUsize>,
    handle: reactor::Handle,
    join: Option<thread::JoinHandle<()>>,
}

impl Drop for IOContextShutdown {
    fn drop(&mut self) {
        use tokio_executor::park::Unpark;

        gst_debug!(CONTEXT_CAT, "Shutting down reactor thread '{}'", self.name);
        self.shutdown.store(SHUTDOWN_NOW, atomic::Ordering::SeqCst);
        gst_trace!(CONTEXT_CAT, "Waiting for reactor '{}' shutdown", self.name);
        // After being unparked, the next turn() is guaranteed to finish immediately,
        // as such there is no race condition between checking for shutdown and setting
        // shutdown.
        self.handle.unpark();
        let _ = self.join.take().unwrap().join();
    }
}

#[derive(Clone)]
pub struct IOContext(Arc<IOContextInner>);

impl glib::subclass::boxed::BoxedType for IOContext {
    const NAME: &'static str = "TsIOContext";

    glib_boxed_type!();
}

glib_boxed_derive_traits!(IOContext);

struct IOContextInner {
    name: String,
    runtime_handle: Mutex<tokio_current_thread::Handle>,
    reactor_handle: reactor::Handle,
    timers: Arc<Mutex<BinaryHeap<TimerEntry>>>,
    // Only used for dropping
    _shutdown: IOContextShutdown,
    pending_futures: Mutex<(
        u64,
        HashMap<u64, FuturesUnordered<Box<dyn Future<Item = (), Error = ()> + Send + 'static>>>,
    )>,
}

impl Drop for IOContextInner {
    fn drop(&mut self) {
        let mut contexts = CONTEXTS.lock().unwrap();
        gst_debug!(CONTEXT_CAT, "Finalizing context '{}'", self.name);
        contexts.remove(&self.name);
    }
}

impl IOContext {
    pub fn new(name: &str, wait: u32) -> Result<Self, io::Error> {
        let mut contexts = CONTEXTS.lock().unwrap();
        if let Some(context) = contexts.get(name) {
            if let Some(context) = context.upgrade() {
                gst_debug!(CONTEXT_CAT, "Reusing existing context '{}'", name);
                return Ok(IOContext(context));
            }
        }

        let reactor = reactor::Reactor::new()?;
        let reactor_handle = reactor.handle().clone();

        let timers = Arc::new(Mutex::new(BinaryHeap::new()));

        let (runtime_handle, shutdown) =
            IOContextRunner::start(name, wait, reactor, timers.clone());

        let context = Arc::new(IOContextInner {
            name: name.into(),
            runtime_handle: Mutex::new(runtime_handle),
            reactor_handle,
            timers,
            _shutdown: shutdown,
            pending_futures: Mutex::new((0, HashMap::new())),
        });
        contexts.insert(name.into(), Arc::downgrade(&context));

        gst_debug!(CONTEXT_CAT, "Created new context '{}'", name);
        Ok(IOContext(context))
    }

    pub fn spawn<F>(&self, future: F)
    where
        F: Future<Item = (), Error = ()> + Send + 'static,
    {
        self.0.runtime_handle.lock().unwrap().spawn(future).unwrap();
    }

    pub fn reactor_handle(&self) -> &reactor::Handle {
        &self.0.reactor_handle
    }

    pub fn acquire_pending_future_id(&self) -> PendingFutureId {
        let mut pending_futures = self.0.pending_futures.lock().unwrap();
        let id = pending_futures.0;
        pending_futures.0 += 1;
        pending_futures.1.insert(id, FuturesUnordered::new());

        PendingFutureId(id)
    }

    pub fn release_pending_future_id(&self, id: PendingFutureId) {
        let mut pending_futures = self.0.pending_futures.lock().unwrap();
        if let Some(fs) = pending_futures.1.remove(&id.0) {
            self.spawn(fs.for_each(|_| Ok(())));
        }
    }

    pub fn add_pending_future<F>(&self, id: PendingFutureId, future: F)
    where
        F: Future<Item = (), Error = ()> + Send + 'static,
    {
        let mut pending_futures = self.0.pending_futures.lock().unwrap();
        let fs = pending_futures.1.get_mut(&id.0).unwrap();
        fs.push(Box::new(future))
    }

    pub fn drain_pending_futures<E: Send + 'static>(
        &self,
        id: PendingFutureId,
    ) -> (Option<oneshot::Sender<()>>, PendingFuturesFuture<E>) {
        let mut pending_futures = self.0.pending_futures.lock().unwrap();
        let fs = pending_futures.1.get_mut(&id.0).unwrap();

        let pending_futures = mem::replace(fs, FuturesUnordered::new());

        if !pending_futures.is_empty() {
            gst_log!(
                CONTEXT_CAT,
                "Scheduling {} pending futures for context '{}' with pending future id {:?}",
                pending_futures.len(),
                self.0.name,
                id,
            );

            let (sender, receiver) = oneshot::channel();

            let future = pending_futures
                .for_each(|_| Ok(()))
                .select(receiver.then(|_| Ok(())))
                .then(|_| Ok(()));

            (Some(sender), future::Either::A(Box::new(future)))
        } else {
            (None, future::Either::B(future::ok(())))
        }
    }
}

pub type PendingFuturesFuture<E> = future::Either<
    Box<dyn Future<Item = (), Error = E> + Send + 'static>,
    future::FutureResult<(), E>,
>;

#[derive(Clone, Copy, Eq, PartialEq, Hash, Debug)]
pub struct PendingFutureId(u64);

impl glib::subclass::boxed::BoxedType for PendingFutureId {
    const NAME: &'static str = "TsPendingFutureId";

    glib_boxed_type!();
}

glib_boxed_derive_traits!(PendingFutureId);

static TIMER_ENTRY_ID: atomic::AtomicUsize = atomic::AtomicUsize::new(0);

// Ad-hoc interval timer implementation for our throttled event loop above
pub struct TimerEntry {
    time: time::Instant,
    id: usize, // for producing a total order
    interval: Option<time::Duration>,
    sender: futures_mpsc::UnboundedSender<()>,
}

impl PartialEq for TimerEntry {
    fn eq(&self, other: &Self) -> bool {
        self.time.eq(&other.time) && self.id.eq(&other.id)
    }
}

impl Eq for TimerEntry {}

impl PartialOrd for TimerEntry {
    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
        Some(self.cmp(&other))
    }
}

impl Ord for TimerEntry {
    fn cmp(&self, other: &Self) -> cmp::Ordering {
        other
            .time
            .cmp(&self.time)
            .then_with(|| other.id.cmp(&self.id))
    }
}

pub struct Interval {
    receiver: futures_mpsc::UnboundedReceiver<()>,
}

impl Interval {
    pub fn new(context: &IOContext, interval: time::Duration) -> Self {
        let (sender, receiver) = futures_mpsc::unbounded();

        let mut timers = context.0.timers.lock().unwrap();
        let entry = TimerEntry {
            time: time::Instant::now(),
            id: TIMER_ENTRY_ID.fetch_add(1, atomic::Ordering::Relaxed),
            interval: Some(interval),
            sender,
        };
        timers.push(entry);

        Self { receiver }
    }
}

impl Stream for Interval {
    type Item = ();
    type Error = ();

    fn poll(&mut self) -> futures::Poll<Option<Self::Item>, Self::Error> {
        self.receiver.poll()
    }
}
<|endoftext|>


<commit_before>use std::{
    future::Future,
    pin::Pin,
    sync::{
        atomic::{AtomicBool, Ordering::Relaxed},
        mpsc::{sync_channel, Receiver, SyncSender, TryRecvError},
    },
    task::{Context, Poll, Waker},
    time::{Duration, Instant},
};

use crate::*;

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

/// An event that happened to a key that a subscriber is interested in.
#[derive(Debug, Clone)]
pub struct Event {
    /// A map of batches for each tree written to in a transaction,
    /// only one of which will be the one subscribed to.
    pub(crate) batches: Arc<[(Tree, Batch)]>,
}

impl Event {
    pub(crate) fn single_update(
        tree: Tree,
        key: IVec,
        value: Option<IVec>,
    ) -> Event {
        Event::single_batch(
            tree,
            Batch { writes: vec![(key, value)].into_iter().collect() },
        )
    }

    pub(crate) fn single_batch(tree: Tree, batch: Batch) -> Event {
        Event::from_batches(vec![(tree, batch)])
    }

    pub(crate) fn from_batches(batches: Vec<(Tree, Batch)>) -> Event {
        Event { batches: Arc::from(batches.into_boxed_slice()) }
    }

    /// Iterate over each Tree, key, and optional value in this `Event`
    pub fn iter<'a>(
        &'a self,
    ) -> Box<dyn 'a + Iterator<Item = (&'a Tree, &'a IVec, &'a Option<IVec>)>>
    {
        self.into_iter()
    }
}

impl<'a> IntoIterator for &'a Event {
    type Item = (&'a Tree, &'a IVec, &'a Option<IVec>);
    type IntoIter = Box<dyn 'a + Iterator<Item = Self::Item>>;

    fn into_iter(self) -> Self::IntoIter {
        Box::new(self.batches.iter().flat_map(|(ref tree, ref batch)| {
            batch.writes.iter().map(move |(k, v_opt)| (tree, k, v_opt))
        }))
    }
}

type Senders = Map<usize, (Option<Waker>, SyncSender<OneShot<Option<Event>>>)>;

/// A subscriber listening on a specified prefix
///
/// `Subscriber` implements both `Iterator<Item = Event>`
/// and `Future<Output=Option<Event>>`
///
/// # Examples
///
/// Synchronous, blocking subscriber:
/// ```
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
/// use sled::{Config, Event};
/// let config = Config::new().temporary(true);
///
/// let tree = config.open()?;
///
/// // watch all events by subscribing to the empty prefix
/// let mut subscriber = tree.watch_prefix(vec![]);
///
/// let tree_2 = tree.clone();
/// let thread = std::thread::spawn(move || {
///     tree.insert(vec![0], vec![1])
/// });
///
/// // `Subscription` implements `Iterator<Item=Event>`
/// for event in subscriber.take(1) {
///     // Events occur due to single key operations,
///     // batches, or transactions. The tree is included
///     // so that you may perform a new transaction or
///     // operation in response to the event.
///     for (tree, key, value_opt) in &event {
///         if let Some(value) = value_opt {
///             // key `key` was set to value `value`
///         } else {
///             // key `key` was removed
///         }
///     }
/// }
///
/// # thread.join().unwrap();
/// # Ok(())
/// # }
/// ```
/// Aynchronous, non-blocking subscriber:
///
/// `Subscription` implements `Future<Output=Option<Event>>`.
///
/// `while let Some(event) = (&mut subscriber).await { /* use it */ }`
pub struct Subscriber {
    id: usize,
    rx: Receiver<OneShot<Option<Event>>>,
    existing: Option<OneShot<Option<Event>>>,
    home: Arc<RwLock<Senders>>,
}

impl Drop for Subscriber {
    fn drop(&mut self) {
        let mut w_senders = self.home.write();
        w_senders.remove(&self.id);
    }
}

impl Subscriber {
    /// Attempts to wait for a value on this `Subscriber`, returning
    /// an error if no event arrives within the provided `Duration`
    /// or if the backing `Db` shuts down.
    pub fn next_timeout(
        &self,
        mut timeout: Duration,
    ) -> std::result::Result<Event, std::sync::mpsc::RecvTimeoutError> {
        loop {
            let start = Instant::now();
            let future_rx = self.rx.recv_timeout(timeout)?;
            timeout =
                if let Some(timeout) = timeout.checked_sub(start.elapsed()) {
                    timeout
                } else {
                    Duration::from_nanos(0)
                };

            let start = Instant::now();
            if let Some(event) = future_rx.wait_timeout(timeout)? {
                return Ok(event);
            }
            timeout =
                if let Some(timeout) = timeout.checked_sub(start.elapsed()) {
                    timeout
                } else {
                    Duration::from_nanos(0)
                };
        }
    }

    fn poll_oneshot(
        self: &mut Pin<&mut Self>,
        mut oneshot: OneShot<Option<Event>>,
        cx: &mut Context<'_>,
    ) -> Option<Poll<Option<Event>>> {
        match Future::poll(Pin::new(&mut oneshot), cx) {
            Poll::Ready(Some(event)) => Some(Poll::Ready(event)),
            Poll::Ready(None) => None,
            Poll::Pending => {
                self.existing = Some(oneshot);
                Some(Poll::Pending)
            }
        }
    }
}

impl Future for Subscriber {
    type Output = Option<Event>;

    fn poll(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Self::Output> {
        loop {
            if let Some(future_rx) = self.existing.take() {
                if let Some(poll) = self.poll_oneshot(future_rx, cx) {
                    return poll;
                }
            }

            match self.rx.try_recv() {
                Ok(future_rx) => {
                    if let Some(poll) = self.poll_oneshot(future_rx, cx) {
                        return poll;
                    }
                }
                Err(TryRecvError::Empty) => break,
                Err(TryRecvError::Disconnected) => return Poll::Ready(None),
            }
        }
        let mut home = self.home.write();
        let entry = home.get_mut(&self.id).unwrap();
        entry.0 = Some(cx.waker().clone());
        Poll::Pending
    }
}

impl Iterator for Subscriber {
    type Item = Event;

    fn next(&mut self) -> Option<Event> {
        loop {
            let future_rx = self.rx.recv().ok()?;
            match future_rx.wait() {
                Some(Some(event)) => return Some(event),
                Some(None) => return None,
                None => continue,
            }
        }
    }
}

#[derive(Debug, Default)]
pub(crate) struct Subscribers {
    watched: RwLock<BTreeMap<Vec<u8>, Arc<RwLock<Senders>>>>,
    ever_used: AtomicBool,
}

impl Drop for Subscribers {
    fn drop(&mut self) {
        let watched = self.watched.read();

        for senders in watched.values() {
            let senders = std::mem::take(&mut *senders.write());
            for (_, (waker, sender)) in senders {
                drop(sender);
                if let Some(waker) = waker {
                    waker.wake();
                }
            }
        }
    }
}

impl Subscribers {
    pub(crate) fn register(&self, prefix: &[u8]) -> Subscriber {
        self.ever_used.store(true, Relaxed);
        let r_mu = {
            let r_mu = self.watched.read();
            if r_mu.contains_key(prefix) {
                r_mu
            } else {
                drop(r_mu);
                let mut w_mu = self.watched.write();
                if !w_mu.contains_key(prefix) {
                    let old = w_mu.insert(
                        prefix.to_vec(),
                        Arc::new(RwLock::new(Map::default())),
                    );
                    assert!(old.is_none());
                }
                drop(w_mu);
                self.watched.read()
            }
        };

        let (tx, rx) = sync_channel(1024);

        let arc_senders = &r_mu[prefix];
        let mut w_senders = arc_senders.write();

        let id = ID_GEN.fetch_add(1, Relaxed);

        w_senders.insert(id, (None, tx));

        Subscriber { id, rx, existing: None, home: arc_senders.clone() }
    }

    pub(crate) fn reserve_batch(
        &self,
        batch: &Batch,
    ) -> Option<ReservedBroadcast> {
        if !self.ever_used.load(Relaxed) {
            return None;
        }

        let r_mu = self.watched.read();

        let mut skip_indices = std::collections::HashSet::new();
        let mut subscribers = vec![];

        for key in batch.writes.keys() {
            for (idx, (prefix, subs_rwl)) in r_mu.iter().enumerate() {
                if key.starts_with(prefix) && !skip_indices.contains(&idx) {
                    skip_indices.insert(idx);
                    let subs = subs_rwl.read();

                    for (_id, (waker, sender)) in subs.iter() {
                        let (tx, rx) = OneShot::pair();
                        if sender.send(rx).is_err() {
                            continue;
                        }
                        subscribers.push((waker.clone(), tx));
                    }
                }
            }
        }

        if subscribers.is_empty() {
            None
        } else {
            Some(ReservedBroadcast { subscribers })
        }
    }

    pub(crate) fn reserve<R: AsRef<[u8]>>(
        &self,
        key: R,
    ) -> Option<ReservedBroadcast> {
        if !self.ever_used.load(Relaxed) {
            return None;
        }

        let r_mu = self.watched.read();
        let prefixes = r_mu.iter().filter(|(k, _)| key.as_ref().starts_with(k));

        let mut subscribers = vec![];

        for (_, subs_rwl) in prefixes {
            let subs = subs_rwl.read();

            for (_id, (waker, sender)) in subs.iter() {
                let (tx, rx) = OneShot::pair();
                if sender.send(rx).is_err() {
                    continue;
                }
                subscribers.push((waker.clone(), tx));
            }
        }

        if subscribers.is_empty() {
            None
        } else {
            Some(ReservedBroadcast { subscribers })
        }
    }
}

pub(crate) struct ReservedBroadcast {
    subscribers: Vec<(Option<Waker>, OneShotFiller<Option<Event>>)>,
}

impl ReservedBroadcast {
    pub fn complete(self, event: &Event) {
        let iter = self.subscribers.into_iter();

        for (waker, tx) in iter {
            tx.fill(Some(event.clone()));
            if let Some(waker) = waker {
                waker.wake();
            }
        }
    }
}
<commit_msg>Follow @rrichardson's lead and only poll oneshots in one place<commit_after>use std::{
    future::Future,
    pin::Pin,
    sync::{
        atomic::{AtomicBool, Ordering::Relaxed},
        mpsc::{sync_channel, Receiver, SyncSender, TryRecvError},
    },
    task::{Context, Poll, Waker},
    time::{Duration, Instant},
};

use crate::*;

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

/// An event that happened to a key that a subscriber is interested in.
#[derive(Debug, Clone)]
pub struct Event {
    /// A map of batches for each tree written to in a transaction,
    /// only one of which will be the one subscribed to.
    pub(crate) batches: Arc<[(Tree, Batch)]>,
}

impl Event {
    pub(crate) fn single_update(
        tree: Tree,
        key: IVec,
        value: Option<IVec>,
    ) -> Event {
        Event::single_batch(
            tree,
            Batch { writes: vec![(key, value)].into_iter().collect() },
        )
    }

    pub(crate) fn single_batch(tree: Tree, batch: Batch) -> Event {
        Event::from_batches(vec![(tree, batch)])
    }

    pub(crate) fn from_batches(batches: Vec<(Tree, Batch)>) -> Event {
        Event { batches: Arc::from(batches.into_boxed_slice()) }
    }

    /// Iterate over each Tree, key, and optional value in this `Event`
    pub fn iter<'a>(
        &'a self,
    ) -> Box<dyn 'a + Iterator<Item = (&'a Tree, &'a IVec, &'a Option<IVec>)>>
    {
        self.into_iter()
    }
}

impl<'a> IntoIterator for &'a Event {
    type Item = (&'a Tree, &'a IVec, &'a Option<IVec>);
    type IntoIter = Box<dyn 'a + Iterator<Item = Self::Item>>;

    fn into_iter(self) -> Self::IntoIter {
        Box::new(self.batches.iter().flat_map(|(ref tree, ref batch)| {
            batch.writes.iter().map(move |(k, v_opt)| (tree, k, v_opt))
        }))
    }
}

type Senders = Map<usize, (Option<Waker>, SyncSender<OneShot<Option<Event>>>)>;

/// A subscriber listening on a specified prefix
///
/// `Subscriber` implements both `Iterator<Item = Event>`
/// and `Future<Output=Option<Event>>`
///
/// # Examples
///
/// Synchronous, blocking subscriber:
/// ```
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
/// use sled::{Config, Event};
/// let config = Config::new().temporary(true);
///
/// let tree = config.open()?;
///
/// // watch all events by subscribing to the empty prefix
/// let mut subscriber = tree.watch_prefix(vec![]);
///
/// let tree_2 = tree.clone();
/// let thread = std::thread::spawn(move || {
///     tree.insert(vec![0], vec![1])
/// });
///
/// // `Subscription` implements `Iterator<Item=Event>`
/// for event in subscriber.take(1) {
///     // Events occur due to single key operations,
///     // batches, or transactions. The tree is included
///     // so that you may perform a new transaction or
///     // operation in response to the event.
///     for (tree, key, value_opt) in &event {
///         if let Some(value) = value_opt {
///             // key `key` was set to value `value`
///         } else {
///             // key `key` was removed
///         }
///     }
/// }
///
/// # thread.join().unwrap();
/// # Ok(())
/// # }
/// ```
/// Aynchronous, non-blocking subscriber:
///
/// `Subscription` implements `Future<Output=Option<Event>>`.
///
/// `while let Some(event) = (&mut subscriber).await { /* use it */ }`
pub struct Subscriber {
    id: usize,
    rx: Receiver<OneShot<Option<Event>>>,
    existing: Option<OneShot<Option<Event>>>,
    home: Arc<RwLock<Senders>>,
}

impl Drop for Subscriber {
    fn drop(&mut self) {
        let mut w_senders = self.home.write();
        w_senders.remove(&self.id);
    }
}

impl Subscriber {
    /// Attempts to wait for a value on this `Subscriber`, returning
    /// an error if no event arrives within the provided `Duration`
    /// or if the backing `Db` shuts down.
    pub fn next_timeout(
        &self,
        mut timeout: Duration,
    ) -> std::result::Result<Event, std::sync::mpsc::RecvTimeoutError> {
        loop {
            let start = Instant::now();
            let future_rx = self.rx.recv_timeout(timeout)?;
            timeout =
                if let Some(timeout) = timeout.checked_sub(start.elapsed()) {
                    timeout
                } else {
                    Duration::from_nanos(0)
                };

            let start = Instant::now();
            if let Some(event) = future_rx.wait_timeout(timeout)? {
                return Ok(event);
            }
            timeout =
                if let Some(timeout) = timeout.checked_sub(start.elapsed()) {
                    timeout
                } else {
                    Duration::from_nanos(0)
                };
        }
    }
}

impl Future for Subscriber {
    type Output = Option<Event>;

    fn poll(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Self::Output> {
        loop {
            let mut future_rx = if let Some(future_rx) = self.existing.take() {
                future_rx
            } else {
                match self.rx.try_recv() {
                    Ok(future_rx) => future_rx,
                    Err(TryRecvError::Empty) => break,
                    Err(TryRecvError::Disconnected) => {
                        return Poll::Ready(None)
                    }
                }
            };

            match Future::poll(Pin::new(&mut future_rx), cx) {
                Poll::Ready(Some(event)) => return Poll::Ready(event),
                Poll::Ready(None) => continue,
                Poll::Pending => {
                    self.existing = Some(future_rx);
                    return Poll::Pending;
                }
            }
        }
        let mut home = self.home.write();
        let entry = home.get_mut(&self.id).unwrap();
        entry.0 = Some(cx.waker().clone());
        Poll::Pending
    }
}

impl Iterator for Subscriber {
    type Item = Event;

    fn next(&mut self) -> Option<Event> {
        loop {
            let future_rx = self.rx.recv().ok()?;
            match future_rx.wait() {
                Some(Some(event)) => return Some(event),
                Some(None) => return None,
                None => continue,
            }
        }
    }
}

#[derive(Debug, Default)]
pub(crate) struct Subscribers {
    watched: RwLock<BTreeMap<Vec<u8>, Arc<RwLock<Senders>>>>,
    ever_used: AtomicBool,
}

impl Drop for Subscribers {
    fn drop(&mut self) {
        let watched = self.watched.read();

        for senders in watched.values() {
            let senders = std::mem::take(&mut *senders.write());
            for (_, (waker, sender)) in senders {
                drop(sender);
                if let Some(waker) = waker {
                    waker.wake();
                }
            }
        }
    }
}

impl Subscribers {
    pub(crate) fn register(&self, prefix: &[u8]) -> Subscriber {
        self.ever_used.store(true, Relaxed);
        let r_mu = {
            let r_mu = self.watched.read();
            if r_mu.contains_key(prefix) {
                r_mu
            } else {
                drop(r_mu);
                let mut w_mu = self.watched.write();
                if !w_mu.contains_key(prefix) {
                    let old = w_mu.insert(
                        prefix.to_vec(),
                        Arc::new(RwLock::new(Map::default())),
                    );
                    assert!(old.is_none());
                }
                drop(w_mu);
                self.watched.read()
            }
        };

        let (tx, rx) = sync_channel(1024);

        let arc_senders = &r_mu[prefix];
        let mut w_senders = arc_senders.write();

        let id = ID_GEN.fetch_add(1, Relaxed);

        w_senders.insert(id, (None, tx));

        Subscriber { id, rx, existing: None, home: arc_senders.clone() }
    }

    pub(crate) fn reserve_batch(
        &self,
        batch: &Batch,
    ) -> Option<ReservedBroadcast> {
        if !self.ever_used.load(Relaxed) {
            return None;
        }

        let r_mu = self.watched.read();

        let mut skip_indices = std::collections::HashSet::new();
        let mut subscribers = vec![];

        for key in batch.writes.keys() {
            for (idx, (prefix, subs_rwl)) in r_mu.iter().enumerate() {
                if key.starts_with(prefix) && !skip_indices.contains(&idx) {
                    skip_indices.insert(idx);
                    let subs = subs_rwl.read();

                    for (_id, (waker, sender)) in subs.iter() {
                        let (tx, rx) = OneShot::pair();
                        if sender.send(rx).is_err() {
                            continue;
                        }
                        subscribers.push((waker.clone(), tx));
                    }
                }
            }
        }

        if subscribers.is_empty() {
            None
        } else {
            Some(ReservedBroadcast { subscribers })
        }
    }

    pub(crate) fn reserve<R: AsRef<[u8]>>(
        &self,
        key: R,
    ) -> Option<ReservedBroadcast> {
        if !self.ever_used.load(Relaxed) {
            return None;
        }

        let r_mu = self.watched.read();
        let prefixes = r_mu.iter().filter(|(k, _)| key.as_ref().starts_with(k));

        let mut subscribers = vec![];

        for (_, subs_rwl) in prefixes {
            let subs = subs_rwl.read();

            for (_id, (waker, sender)) in subs.iter() {
                let (tx, rx) = OneShot::pair();
                if sender.send(rx).is_err() {
                    continue;
                }
                subscribers.push((waker.clone(), tx));
            }
        }

        if subscribers.is_empty() {
            None
        } else {
            Some(ReservedBroadcast { subscribers })
        }
    }
}

pub(crate) struct ReservedBroadcast {
    subscribers: Vec<(Option<Waker>, OneShotFiller<Option<Event>>)>,
}

impl ReservedBroadcast {
    pub fn complete(self, event: &Event) {
        let iter = self.subscribers.into_iter();

        for (waker, tx) in iter {
            tx.fill(Some(event.clone()));
            if let Some(waker) = waker {
                waker.wake();
            }
        }
    }
}
<|endoftext|>


<commit_before>#[derive(Clone)]
pub enum Yaml {
    String(String),
    Int(i64),
    List(Vec<Yaml>),
    Map(Vec<(String, Yaml)>),
    Bool(bool),
}

impl Yaml {
    pub fn to_yaml_rust(&self) -> yaml_rust::Yaml {
        match self {
            Yaml::String(s) => yaml_rust::Yaml::String(s.clone()),
            Yaml::List(l) => yaml_rust::Yaml::Array(l.iter().map(|i| i.to_yaml_rust()).collect()),
            Yaml::Map(m) => yaml_rust::Yaml::Hash(
                m.iter()
                    .map(|(k, v)| (yaml_rust::Yaml::String(k.clone()), v.to_yaml_rust()))
                    .collect(),
            ),
            Yaml::Int(i) => yaml_rust::Yaml::Integer(*i),
            Yaml::Bool(b) => yaml_rust::Yaml::Boolean(*b),
        }
    }
}

impl From<&Yaml> for Yaml {
    fn from(y: &Yaml) -> Self {
        y.clone()
    }
}

impl From<String> for Yaml {
    fn from(s: String) -> Self {
        Yaml::String(s)
    }
}

impl From<&str> for Yaml {
    fn from(s: &str) -> Self {
        Yaml::String(s.to_owned())
    }
}

impl From<&&str> for Yaml {
    fn from(s: &&str) -> Self {
        Yaml::String((*s).to_owned())
    }
}

impl From<bool> for Yaml {
    fn from(b: bool) -> Self {
        Yaml::Bool(b)
    }
}

impl<T: Into<Yaml>> From<Vec<T>> for Yaml {
    fn from(v: Vec<T>) -> Self {
        Yaml::List(v.into_iter().map(|t| t.into()).collect())
    }
}

impl Yaml {
    pub fn map<K: Into<String>, V: Into<Yaml>, E: IntoIterator<Item = (K, V)>>(entries: E) -> Yaml {
        Yaml::Map(
            entries
                .into_iter()
                .map(|(k, v)| (k.into(), v.into()))
                .collect(),
        )
    }

    pub fn list<V: Into<Yaml>, E: IntoIterator<Item = V>>(values: E) -> Yaml {
        Yaml::List(values.into_iter().map(|v| v.into()).collect())
    }

    pub fn string<S: Into<String>>(s: S) -> Yaml {
        Yaml::String(s.into())
    }
}

#[derive(Default)]
pub struct YamlWriter {
    pub buffer: String,
}

impl YamlWriter {
    pub fn write_line(&mut self, line: &str) {
        self.buffer.push_str(line);
        self.buffer.push_str("\n");
    }

    pub fn write_yaml(&mut self, yaml: &Yaml) {
        yaml_rust::emitter::YamlEmitter::new(&mut self.buffer)
            .compact()
            .dump(&yaml.to_yaml_rust())
            .unwrap();
    }
}
<commit_msg>Fix compilation<commit_after>#[derive(Clone)]
pub enum Yaml {
    String(String),
    Int(i64),
    List(Vec<Yaml>),
    Map(Vec<(String, Yaml)>),
    Bool(bool),
}

impl Yaml {
    pub fn to_yaml_rust(&self) -> yaml_rust::Yaml {
        match self {
            Yaml::String(s) => yaml_rust::Yaml::String(s.clone()),
            Yaml::List(l) => yaml_rust::Yaml::Array(l.iter().map(|i| i.to_yaml_rust()).collect()),
            Yaml::Map(m) => yaml_rust::Yaml::Hash(
                m.iter()
                    .map(|(k, v)| (yaml_rust::Yaml::String(k.clone()), v.to_yaml_rust()))
                    .collect(),
            ),
            Yaml::Int(i) => yaml_rust::Yaml::Integer(*i),
            Yaml::Bool(b) => yaml_rust::Yaml::Boolean(*b),
        }
    }
}

impl From<&Yaml> for Yaml {
    fn from(y: &Yaml) -> Self {
        y.clone()
    }
}

impl From<String> for Yaml {
    fn from(s: String) -> Self {
        Yaml::String(s)
    }
}

impl From<&str> for Yaml {
    fn from(s: &str) -> Self {
        Yaml::String(s.to_owned())
    }
}

impl From<&&str> for Yaml {
    fn from(s: &&str) -> Self {
        Yaml::String((*s).to_owned())
    }
}

impl From<bool> for Yaml {
    fn from(b: bool) -> Self {
        Yaml::Bool(b)
    }
}

impl<T: Into<Yaml>> From<Vec<T>> for Yaml {
    fn from(v: Vec<T>) -> Self {
        Yaml::List(v.into_iter().map(|t| t.into()).collect())
    }
}

impl Yaml {
    pub fn map<K: Into<String>, V: Into<Yaml>, E: IntoIterator<Item = (K, V)>>(entries: E) -> Yaml {
        Yaml::Map(
            entries
                .into_iter()
                .map(|(k, v)| (k.into(), v.into()))
                .collect(),
        )
    }

    pub fn list<V: Into<Yaml>, E: IntoIterator<Item = V>>(values: E) -> Yaml {
        Yaml::List(values.into_iter().map(|v| v.into()).collect())
    }

    pub fn string<S: Into<String>>(s: S) -> Yaml {
        Yaml::String(s.into())
    }
}

#[derive(Default)]
pub struct YamlWriter {
    pub buffer: String,
}

impl YamlWriter {
    pub fn write_line(&mut self, line: &str) {
        self.buffer.push_str(line);
        self.buffer.push_str("\n");
    }

    pub fn write_yaml(&mut self, yaml: &Yaml) {
        yaml_rust::emitter::YamlEmitter::new(&mut self.buffer)
            .dump(&yaml.to_yaml_rust())
            .unwrap();
    }
}
<|endoftext|>


<commit_before>use std::str;
use types::{Frequency, Field, Message, Enumerator, MessageOrEnum, FileDescriptor, Syntax};
use nom::{multispace, digit};

fn is_word(b: u8) -> bool {
    match b {
        b'a'...b'z' | b'A'...b'Z' | b'0'...b'9' | b'_' => true,
        _ => false
    }
}

named!(word<&str>, map_res!(take_while!(is_word), str::from_utf8));

named!(comment<()>, do_parse!(tag!("//") >> take_until_and_consume!("\n") >> ()));
named!(block_comment<()>, do_parse!(tag!("/*") >> take_until_and_consume!("*/") >> ()));

/// word break: multispace or comment
named!(br<()>, alt!(map!(multispace, |_| ()) | comment | block_comment));

named!(syntax<Syntax>, do_parse!(tag!("syntax") >> many0!(br) >> tag!("=") >>
    proto: alt!(tag!("\"proto2\"") => { |_| Syntax::Proto2 } |
                tag!("\"proto3\"") => { |_| Syntax::Proto3 }) >> (proto)));

named!(default_value<&str>, do_parse!(
    tag!("[") >> many0!(br) >> tag!("default") >> many0!(br) >> tag!("=") >> many0!(br) >>
    default: word >> many0!(br) >> tag!("]") >>
    (default)));

named!(deprecated<bool>, do_parse!(
    tag!("[") >> many0!(br) >> tag!("deprecated") >> many0!(br) >> tag!("=") >> many0!(br) >>
    deprecated: map_res!(word, str::FromStr::from_str) >> many0!(br) >> tag!("]") >>
    (deprecated)));

named!(packed<bool>, do_parse!(
    tag!("[") >> many0!(br) >> tag!("packed") >> many0!(br) >> tag!("=") >> many0!(br) >>
    packed: map_res!(word, str::FromStr::from_str) >> many0!(br) >> tag!("]") >>
    (packed)));

named!(frequency<Frequency>,
       alt!(tag!("optional") => { |_| Frequency::Optional } |
            tag!("repeated") => { |_| Frequency::Repeated } |
            tag!("required") => { |_| Frequency::Required } ));

named!(message_field<Field>, do_parse!(
    frequency: opt!(frequency) >> many1!(br) >>
    typ: word >> many1!(br) >>
    name: word >> many0!(br) >>
    tag!("=") >> many0!(br) >>
    number: map_res!(map_res!(digit, str::from_utf8), str::FromStr::from_str) >> many0!(br) >>
    default: opt!(default_value) >> many0!(br) >>
    deprecated: opt!(deprecated) >> many0!(br) >>
    packed: opt!(packed) >> many0!(br) >> tag!(";") >> many0!(br) >>
    (Field {
       name: name,
       frequency: frequency.unwrap_or(Frequency::Optional),
       typ: typ,
       number: number,
       default: default,
       packed: packed,
       boxed: false,
       deprecated: deprecated.unwrap_or(false),
    })));

named!(message<Message>, do_parse!(
    tag!("message") >> many0!(br) >>
    name: word >> many0!(br) >>
    tag!("{") >> many0!(br) >>
    fields: many0!(message_field) >>
    tag!("}") >> many0!(br) >>
    (Message { name: name, fields: fields })));

named!(enum_field<(&str, i32)>, do_parse!(
    name: word >> many0!(br) >>
    tag!("=") >> many0!(br) >>
    number: map_res!(map_res!(digit, str::from_utf8), str::FromStr::from_str) >> many0!(br) >>
    tag!(";") >> many0!(br) >>
    ((name, number))));

named!(enumerator<Enumerator>, do_parse!(
    tag!("enum") >> many1!(br) >>
    name: word >> many0!(br) >>
    tag!("{") >> many0!(br) >>
    fields: many0!(enum_field) >>
    tag!("}") >> many0!(br) >>
    (Enumerator { name: name, fields: fields })));

named!(ignore<()>, do_parse!(
    alt!(tag!("package") | tag!("option") | tag!("import")) >> many1!(br) >>
    take_until_and_consume!(";") >> many0!(br) >> ()));

named!(message_or_enum<MessageOrEnum>, alt!(
         message => { |m| MessageOrEnum::Msg(m) } |
         enumerator => { |e| MessageOrEnum::Enum(e) } |
         ignore => { |_| MessageOrEnum::Ignore } ));

named!(pub file_descriptor<FileDescriptor>, do_parse!(
    many0!(br) >> syntax: opt!(syntax) >> many0!(br) >>
    message_and_enums: many0!(message_or_enum) >>
    (FileDescriptor {
        syntax: syntax.unwrap_or(Syntax::Proto2),
        message_and_enums: message_and_enums,
        messages: Vec::new(),
        enums: Vec::new(),
    })));

#[test]
fn test_message() {
    let msg = r#"message ReferenceData
{
    repeated ScenarioInfo  scenarioSet = 1;
    repeated CalculatedObjectInfo calculatedObjectSet = 2;
    repeated RiskFactorList riskFactorListSet = 3;
    repeated RiskMaturityInfo riskMaturitySet = 4;
    repeated IndicatorInfo indicatorSet = 5;
    repeated RiskStrikeInfo riskStrikeSet = 6;
    repeated FreeProjectionList freeProjectionListSet = 7;
    repeated ValidationProperty ValidationSet = 8;
    repeated CalcProperties calcPropertiesSet = 9;
    repeated MaturityInfo maturitySet = 10;
}"#;

    let mess = message(msg.as_bytes());
    if let ::nom::IResult::Done(_, mess) = mess {
        assert_eq!(10, mess.fields.len());
    }
}

#[test]
fn test_enum() {
    let msg = r#"enum PairingStatus {
            DEALPAIRED        = 0;
            INVENTORYORPHAN   = 1;
            CALCULATEDORPHAN  = 2;
            CANCELED          = 3;
}"#;

    let mess = enumerator(msg.as_bytes());
    if let ::nom::IResult::Done(_, mess) = mess {
        assert_eq!(4, mess.fields.len());
    }
}

#[test]
fn test_ignore() {
    let msg = r#"package com.test.v0;

option optimize_for = SPEED;
"#;

    match ignore(msg.as_bytes()) {
        ::nom::IResult::Done(_, _) => (),
        e => panic!("Expecting done {:?}", e),
    }
}
<commit_msg>ignore services<commit_after>use std::str;
use types::{Frequency, Field, Message, Enumerator, MessageOrEnum, FileDescriptor, Syntax};
use nom::{multispace, digit};

fn is_word(b: u8) -> bool {
    match b {
        b'a'...b'z' | b'A'...b'Z' | b'0'...b'9' | b'_' => true,
        _ => false
    }
}

named!(word<&str>, map_res!(take_while!(is_word), str::from_utf8));

named!(comment<()>, do_parse!(tag!("//") >> take_until_and_consume!("\n") >> ()));
named!(block_comment<()>, do_parse!(tag!("/*") >> take_until_and_consume!("*/") >> ()));

/// word break: multispace or comment
named!(br<()>, alt!(map!(multispace, |_| ()) | comment | block_comment));

named!(syntax<Syntax>, do_parse!(tag!("syntax") >> many0!(br) >> tag!("=") >>
    proto: alt!(tag!("\"proto2\"") => { |_| Syntax::Proto2 } |
                tag!("\"proto3\"") => { |_| Syntax::Proto3 }) >> (proto)));

named!(default_value<&str>, do_parse!(
    tag!("[") >> many0!(br) >> tag!("default") >> many0!(br) >> tag!("=") >> many0!(br) >>
    default: word >> many0!(br) >> tag!("]") >>
    (default)));

named!(deprecated<bool>, do_parse!(
    tag!("[") >> many0!(br) >> tag!("deprecated") >> many0!(br) >> tag!("=") >> many0!(br) >>
    deprecated: map_res!(word, str::FromStr::from_str) >> many0!(br) >> tag!("]") >>
    (deprecated)));

named!(packed<bool>, do_parse!(
    tag!("[") >> many0!(br) >> tag!("packed") >> many0!(br) >> tag!("=") >> many0!(br) >>
    packed: map_res!(word, str::FromStr::from_str) >> many0!(br) >> tag!("]") >>
    (packed)));

named!(frequency<Frequency>,
       alt!(tag!("optional") => { |_| Frequency::Optional } |
            tag!("repeated") => { |_| Frequency::Repeated } |
            tag!("required") => { |_| Frequency::Required } ));

named!(message_field<Field>, do_parse!(
    frequency: opt!(frequency) >> many1!(br) >>
    typ: word >> many1!(br) >>
    name: word >> many0!(br) >>
    tag!("=") >> many0!(br) >>
    number: map_res!(map_res!(digit, str::from_utf8), str::FromStr::from_str) >> many0!(br) >>
    default: opt!(default_value) >> many0!(br) >>
    deprecated: opt!(deprecated) >> many0!(br) >>
    packed: opt!(packed) >> many0!(br) >> tag!(";") >> many0!(br) >>
    (Field {
       name: name,
       frequency: frequency.unwrap_or(Frequency::Optional),
       typ: typ,
       number: number,
       default: default,
       packed: packed,
       boxed: false,
       deprecated: deprecated.unwrap_or(false),
    })));

named!(message<Message>, do_parse!(
    tag!("message") >> many0!(br) >>
    name: word >> many0!(br) >>
    tag!("{") >> many0!(br) >>
    fields: many0!(message_field) >>
    tag!("}") >> many0!(br) >>
    (Message { name: name, fields: fields })));

named!(enum_field<(&str, i32)>, do_parse!(
    name: word >> many0!(br) >>
    tag!("=") >> many0!(br) >>
    number: map_res!(map_res!(digit, str::from_utf8), str::FromStr::from_str) >> many0!(br) >>
    tag!(";") >> many0!(br) >>
    ((name, number))));

named!(enumerator<Enumerator>, do_parse!(
    tag!("enum") >> many1!(br) >>
    name: word >> many0!(br) >>
    tag!("{") >> many0!(br) >>
    fields: many0!(enum_field) >>
    tag!("}") >> many0!(br) >>
    (Enumerator { name: name, fields: fields })));

named!(ignore<()>, do_parse!(
    alt!(tag!("package") | tag!("option") | tag!("import")) >> many1!(br) >>
    take_until_and_consume!(";") >> many0!(br) >> ()));

named!(service_ignore<()>, do_parse!(tag!("service") >> many1!(br) >> word >> many0!(br) >> tag!("{") >>
                                     take_until_and_consume!("}") >> many0!(br) >> ()));

named!(message_or_enum<MessageOrEnum>, alt!(
         message => { |m| MessageOrEnum::Msg(m) } |
         enumerator => { |e| MessageOrEnum::Enum(e) } |
         ignore => { |_| MessageOrEnum::Ignore } |
         service_ignore => { |_| MessageOrEnum::Ignore } ));

named!(pub file_descriptor<FileDescriptor>, do_parse!(
    many0!(br) >> syntax: opt!(syntax) >> many0!(br) >>
    message_and_enums: many0!(message_or_enum) >>
    (FileDescriptor {
        syntax: syntax.unwrap_or(Syntax::Proto2),
        message_and_enums: message_and_enums,
        messages: Vec::new(),
        enums: Vec::new(),
    })));

#[test]
fn test_message() {
    let msg = r#"message ReferenceData
{
    repeated ScenarioInfo  scenarioSet = 1;
    repeated CalculatedObjectInfo calculatedObjectSet = 2;
    repeated RiskFactorList riskFactorListSet = 3;
    repeated RiskMaturityInfo riskMaturitySet = 4;
    repeated IndicatorInfo indicatorSet = 5;
    repeated RiskStrikeInfo riskStrikeSet = 6;
    repeated FreeProjectionList freeProjectionListSet = 7;
    repeated ValidationProperty ValidationSet = 8;
    repeated CalcProperties calcPropertiesSet = 9;
    repeated MaturityInfo maturitySet = 10;
}"#;

    let mess = message(msg.as_bytes());
    if let ::nom::IResult::Done(_, mess) = mess {
        assert_eq!(10, mess.fields.len());
    }
}

#[test]
fn test_enum() {
    let msg = r#"enum PairingStatus {
            DEALPAIRED        = 0;
            INVENTORYORPHAN   = 1;
            CALCULATEDORPHAN  = 2;
            CANCELED          = 3;
}"#;

    let mess = enumerator(msg.as_bytes());
    if let ::nom::IResult::Done(_, mess) = mess {
        assert_eq!(4, mess.fields.len());
    }
}

#[test]
fn test_ignore() {
    let msg = r#"package com.test.v0;

option optimize_for = SPEED;
"#;

    match ignore(msg.as_bytes()) {
        ::nom::IResult::Done(_, _) => (),
        e => panic!("Expecting done {:?}", e),
    }
}
<|endoftext|>