<commit_before>//! Manages serial port state and communication with clients,
//! and handling requests / responses

use std::collections::HashSet;
use std::sync::mpsc::{Receiver, TryRecvError};
use std::thread;
use std::time::Duration;

use base64;

use errors::*;
use messages::*;
use sub_manager::*;
use port_manager::*;
use writelock_manager::*;
use common::*;

/// Serial port management module supporting one
/// writer and multiple readers
///
/// Clients can lock a port for writing, but
/// subscribe to data from multiple ports for reads
///
/// The Manager takes actions in response to
/// [SerialRequest::*](../messages/index.html) messages sent on
/// on its receiver
pub struct Manager {
  /// Manage write lock status
  writelock_manager: WriteLockManager,
  /// Manage ports
  port_manager: PortManager,
  /// Manage subscriptions
  sub_manager: SubscriptionManager,
  /// Receiver for serial requests
  receiver: Receiver<(String, SerialRequest)>,
  /// Receiver for response subscription requests
  subsc_receiver: SubscReceiver,
}

impl Manager {
  ///Constructor
  pub fn new(
    receiver: Receiver<(String, SerialRequest)>,
    subsc_receiver: SubscReceiver,
  ) -> Manager {
    Manager {
      writelock_manager: WriteLockManager::new(),
      port_manager: PortManager::new(),
      sub_manager: SubscriptionManager::new(),
      receiver: receiver,
      subsc_receiver: subsc_receiver,
    }
  }

  ///Spawn an instance in a new thread.
  pub fn spawn(
    receiver: Receiver<(String, SerialRequest)>,
    subsc_receiver: SubscReceiver,
  ) -> thread::JoinHandle<()> {
    thread::spawn(move || { Manager::new(receiver, subsc_receiver).run(); })
  }

  /// Main loop
  fn run(&mut self) {

    // Bad ports we couldn't read from
    // A set of SerialResponse::Errors built from
    // from the serial port read/write error responses
    let mut bad_ports = HashSet::<String>::new();
    // Shutdown flag
    let mut shutdown = false;
    // Check about 30 times a second
    let sleep_dur = Duration::from_millis(33);

    while !shutdown {
      // Sleep for a little bit to avoid pegging cpu
      // TODO Dynamic timing loop
      thread::sleep(sleep_dur);

      // Handle serial operation requests
      match self.receiver.try_recv() {
        Err(e) => {
          match e {
            TryRecvError::Empty => {
              // nothing to do
            }
            TryRecvError::Disconnected => {
              // Remote end hung up, time to shutdown
              info!("Shutting down SerialPortManager");
              shutdown = true;
            }
          }
        }
        Ok(req) => self.handle_serial_request(&req.0, req.1),
      }

      // Check for new data on each port
      for (port_name, result) in self.port_manager.read_all_ports() {
        match result {
          Ok(data) => {
            let response = match String::from_utf8(data) {
              // We need to send as binary
              Err(e) => {
                SerialResponse::Read {
                  port: port_name.to_string(),
                  data: base64::encode(&e.into_bytes()),
                  base64: Some(true),
                }
              }
              Ok(s) => {
                SerialResponse::Read {
                  port: port_name.to_string(),
                  data: s,
                  base64: Some(false),
                }
              }
            };
            self.broadcast_message_for_port(&port_name, response);
          }
          // Send data reads
          Err(e) => {
            warn!("Error reading port!");
            warn!("{}", e);
            bad_ports.insert(port_name);
          }
        }
      }

      //Handle write requests
      let mut recv_count = 0;
      let mut do_recv = true;
      while do_recv {
        recv_count += 1;
        match self.subsc_receiver.try_recv() {
          Ok(sub_request) => self.sub_manager.add_subscription(sub_request),
          Err(e) => {
            match e {
              TryRecvError::Disconnected => {
                // Does this mean all senders have disconnected?
                // Or just one?
                debug!("Got disconnected when trying to get serial request");
              }
              TryRecvError::Empty => do_recv = false,
            }
          }
        }

        // Don't let this loop get stuck here forever
        // starving the outer loop
        if recv_count == 50 {
          do_recv = false;
        }
      }

      // Cleanup bad serial ports that failed read or write
      // We remove them from everything before
      self.cleanup_bad_ports(&bad_ports);
      bad_ports.clear();
    }
  }

  /// Handles and dispatches SerialRequest sent by
  /// the channel
  fn handle_serial_request(&mut self, sub_id: &String, msg: SerialRequest) {
    let response = match msg {
      SerialRequest::Open { port } => self.handle_open_port(sub_id, port),
      SerialRequest::WriteLock { port } => self.handle_write_lock(sub_id, port),
      SerialRequest::ReleaseWriteLock { port } => self.handle_release_write_lock(sub_id, port),
      SerialRequest::Write { port, data, base64 } => {
        self.handle_write_port(sub_id, port, data, base64.unwrap_or(false))
      }
      SerialRequest::Close { port } => self.handle_close_port(sub_id, port),
      SerialRequest::List {} => self.handle_list_ports(sub_id),
    };
    if let Err(e) = response {
      warn!("Error '{}' occured handling serial request message", e);
      // Send error?
      self.send_message(&sub_id, to_serial_response_error(e));

    }
  }

  /// Handle write port requests
  fn handle_write_port(
    &mut self,
    sub_id: &String,
    port_name: String,
    data: String,
    base_64: bool,
  ) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    self.check_owns_writelock(&port_name, &sub_id)?;
    match base_64 {
      true => {
        base64::decode(&data)
          .map_err(|e| ErrorKind::Base64(e).into())
          .and_then(|d| self.port_manager.write_port(&port_name, &d))
          .map(|_| self.send_message(&sub_id, SerialResponse::Wrote { port: port_name }),)
      }
      false => {
        self
          .port_manager
          .write_port(&port_name, data.as_bytes())
          .map(|_| self.send_message(&sub_id, SerialResponse::Wrote { port: port_name }),)
      }
    }
  }

  /// Handle write lock requests
  fn handle_write_lock(&mut self, sub_id: &String, port_name: String) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    self
      .writelock_manager
      .lock_port(&port_name, &sub_id)
      .map(|_| self.send_message(&sub_id, SerialResponse::WriteLocked { port: port_name }),)
  }

  /// Handle write requests
  fn handle_release_write_lock(
    &mut self,
    sub_id: &String,
    port_name: Option<String>,
  ) -> Result<()> {
    self.check_sub_id(sub_id)?;
    match port_name {
      None => {
        self.writelock_manager.unlock_all_ports_for_sub(&sub_id);
        Ok(
          self.send_message(
            &sub_id,
            SerialResponse::WriteLockReleased { port: port_name },
          ),
        )
      }
      Some(port_name) => {
        self
          .writelock_manager
          .unlock_port(&port_name, &sub_id)
          .map(
            |_| {
              self.send_message(
                &sub_id,
                SerialResponse::WriteLockReleased { port: Some(port_name) },
              )
            },
          )
      }
    }
  }

  /// Handle open port requests
  fn handle_open_port(&mut self, sub_id: &String, port_name: String) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    self.port_manager.open_port(&port_name)?;
    self
      .sub_manager
      .add_port(&sub_id, &port_name)
      .map(|_| self.send_message(&sub_id, SerialResponse::Opened { port: port_name }),)
  }

  /// Handle list ports request
  fn handle_list_ports(&mut self, sub_id: &String) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    let port_names: Result<Vec<String>> =
      self
        .port_manager
        .list_ports()
        .map(|v| v.iter().map(|v| v.port_name.clone()).collect());
    let resp = port_names.map(|pns| SerialResponse::List { ports: pns });
    self.send_message(
      &sub_id,
      resp.unwrap_or(
        SerialResponse::Error {
          display: "".to_string(),
          description: "".to_string(),
        },
      ),
    );
    Ok(())
  }

  /// Handle close port requests
  fn handle_close_port(&mut self, sub_id: &String, port_name: Option<String>) -> Result<()> {
    match port_name {
      Some(port_name) => self.handle_close_port_for_sub(sub_id, port_name),
      None => self.handle_close_all_ports_for_sub(sub_id)
    }
  }

  /// Handle closing a signle port for a sub
  fn handle_close_port_for_sub(&mut self, sub_id: &String, port_name: String) -> Result<()> {
    self.sub_manager.remove_port(&sub_id, &port_name)?;
    self.writelock_manager.unlock_port_if_locked_by(&port_name, &sub_id);
    // self.cleanup_ports_with_no_subs();
    let close_resp = SerialResponse::Closed { port: port_name.clone() };
    self.send_message(&sub_id, close_resp);
    Ok(())
  }

  /// Handle closing all ports for sub
  fn handle_close_all_ports_for_sub(&mut self, sub_id: &String) -> Result<()> {
    self.sub_manager.clear_ports(Some(&sub_id));
    self.writelock_manager.unlock_all_ports_for_sub(sub_id);

    // Close ports with no subscribers
    let open_ports = self.port_manager.open_ports();
    let subscribed_ports = self.sub_manager.subscribed_ports();
    let ports_with_no_subs = open_ports.difference(&subscribed_ports);

    // For each open port that isn't subscribed,
    for port_to_close in ports_with_no_subs {
      // close it, REDUNDANT?
      self.port_manager.close_port(&port_to_close);
      // remove the write lock, REDUNDANT?
      self.writelock_manager.clear_lock(&port_to_close);
      // Let them know its closed
      let close_resp = SerialResponse::Closed { port: port_to_close.clone() };
      self.send_message(&sub_id, close_resp);
    }
    Ok(())
  }

  /// Cleanup any bad ports
  fn cleanup_bad_ports(&mut self, bad_ports: &HashSet<String>) {
    for port_name in bad_ports.iter() {
      // Tell everyone port is sick
      let err_resp = to_serial_response_error(ErrorKind::PortReadError(port_name.to_owned()).into(),);
      self.broadcast_message_for_port(port_name, err_resp);
      // Tell everyone the sick ports were closed
      let close_resp = SerialResponse::Closed { port: port_name.clone() };
      self.broadcast_message_for_port(port_name, close_resp);
      // Close bad ports
      self.port_manager.close_port(port_name);
      // Remove write locks on bad ports
      self.writelock_manager.clear_lock(port_name);
      // Remove bad ports from subscriptions
      self.sub_manager.remove_port_from_all(port_name);
    }
  }

  fn send_message(&mut self, sub_id: &String, msg: SerialResponse) {
    if let Err(e) = self.sub_manager.send_message(sub_id, msg) {
      warn!("Error sending serial response to sub_id '{}'", sub_id);
      let mut bad_subs = Vec::new();
      bad_subs.push(e);
      self.cleanup_bad_subs(bad_subs);
    }
  }

  /// Broadcast a message to all subscribers and
  /// then cleanup any subs that errored
  fn broadcast_message(&mut self, msg: SerialResponse) {
    let bad_subs = self.sub_manager.broadcast_message(msg);
    self.cleanup_bad_subs(bad_subs);
  }

  /// Broadcast a message to all subscribers interested in the
  /// given port and then cleanup any subs that errored
  fn broadcast_message_for_port(&mut self, port_name: &String, msg: SerialResponse) {
    let bad_subs = self
      .sub_manager
      .broadcast_message_for_port(port_name, msg);
    self.cleanup_bad_subs(bad_subs);
  }

  /// Cleanup any bad subs where a message send failed
  fn cleanup_bad_subs(&mut self, bad_subs: Vec<Error>) {
    for e in bad_subs {
      if let Error(ErrorKind::SubscriberSendError(sub_id), _) = e {
        // Remove subscriptions
        self.sub_manager.end_subscription(&sub_id);
        // Remove all write locks held by dead subscription
        self.writelock_manager.unlock_all_ports_for_sub(&sub_id);
      }
    }
  }

  fn check_sub_id(&self, sub_id: &String) -> Result<()> {
    self.sub_manager.check_subscription_exists(sub_id)
  }

  fn check_owns_writelock(&self, port_name: &String, sub_id: &String) -> Result<()> {
    self
      .writelock_manager
      .check_owns_write_lock(port_name, sub_id)
  }
}
<commit_msg>Simplify while loop check and exit, and remove unneeded variable<commit_after>//! Manages serial port state and communication with clients,
//! and handling requests / responses

use std::collections::HashSet;
use std::sync::mpsc::{Receiver, TryRecvError};
use std::thread;
use std::time::Duration;

use base64;

use errors::*;
use messages::*;
use sub_manager::*;
use port_manager::*;
use writelock_manager::*;
use common::*;

/// Serial port management module supporting one
/// writer and multiple readers
///
/// Clients can lock a port for writing, but
/// subscribe to data from multiple ports for reads
///
/// The Manager takes actions in response to
/// [SerialRequest::*](../messages/index.html) messages sent on
/// on its receiver
pub struct Manager {
  /// Manage write lock status
  writelock_manager: WriteLockManager,
  /// Manage ports
  port_manager: PortManager,
  /// Manage subscriptions
  sub_manager: SubscriptionManager,
  /// Receiver for serial requests
  receiver: Receiver<(String, SerialRequest)>,
  /// Receiver for response subscription requests
  subsc_receiver: SubscReceiver,
}

impl Manager {
  ///Constructor
  pub fn new(
    receiver: Receiver<(String, SerialRequest)>,
    subsc_receiver: SubscReceiver,
  ) -> Manager {
    Manager {
      writelock_manager: WriteLockManager::new(),
      port_manager: PortManager::new(),
      sub_manager: SubscriptionManager::new(),
      receiver: receiver,
      subsc_receiver: subsc_receiver,
    }
  }

  ///Spawn an instance in a new thread.
  pub fn spawn(
    receiver: Receiver<(String, SerialRequest)>,
    subsc_receiver: SubscReceiver,
  ) -> thread::JoinHandle<()> {
    thread::spawn(move || { Manager::new(receiver, subsc_receiver).run(); })
  }

  /// Main loop
  fn run(&mut self) {

    // Bad ports we couldn't read from
    // A set of SerialResponse::Errors built from
    // from the serial port read/write error responses
    let mut bad_ports = HashSet::<String>::new();
    // Shutdown flag
    let mut shutdown = false;
    // Check about 30 times a second
    let sleep_dur = Duration::from_millis(33);

    while !shutdown {
      // Sleep for a little bit to avoid pegging cpu
      // TODO Dynamic timing loop
      thread::sleep(sleep_dur);

      // Handle serial operation requests
      match self.receiver.try_recv() {
        Err(e) => {
          match e {
            TryRecvError::Empty => {
              // nothing to do
            }
            TryRecvError::Disconnected => {
              // Remote end hung up, time to shutdown
              info!("Shutting down SerialPortManager");
              shutdown = true;
            }
          }
        }
        Ok(req) => self.handle_serial_request(&req.0, req.1),
      }

      // Check for new data on each port
      for (port_name, result) in self.port_manager.read_all_ports() {
        match result {
          Ok(data) => {
            let response = match String::from_utf8(data) {
              // We need to send as binary
              Err(e) => {
                SerialResponse::Read {
                  port: port_name.to_string(),
                  data: base64::encode(&e.into_bytes()),
                  base64: Some(true),
                }
              }
              Ok(s) => {
                SerialResponse::Read {
                  port: port_name.to_string(),
                  data: s,
                  base64: Some(false),
                }
              }
            };
            self.broadcast_message_for_port(&port_name, response);
          }
          // Send data reads
          Err(e) => {
            warn!("Error reading port!");
            warn!("{}", e);
            bad_ports.insert(port_name);
          }
        }
      }

      //Handle write requests
      let mut recv_count = 0;
      while recv_count < 50 {
        recv_count += 1;
        match self.subsc_receiver.try_recv() {
          Ok(sub_request) => self.sub_manager.add_subscription(sub_request),
          Err(e) => {
            match e {
              TryRecvError::Disconnected => {
                // Does this mean all senders have disconnected?
                // Or just one?
                debug!("Got disconnected when trying to get serial request");
              }
              TryRecvError::Empty => break,
            }
          }
        }
      }

      // Cleanup bad serial ports that failed read or write
      // We remove them from everything before
      self.cleanup_bad_ports(&bad_ports);
      bad_ports.clear();
    }
  }

  /// Handles and dispatches SerialRequest sent by
  /// the channel
  fn handle_serial_request(&mut self, sub_id: &String, msg: SerialRequest) {
    let response = match msg {
      SerialRequest::Open { port } => self.handle_open_port(sub_id, port),
      SerialRequest::WriteLock { port } => self.handle_write_lock(sub_id, port),
      SerialRequest::ReleaseWriteLock { port } => self.handle_release_write_lock(sub_id, port),
      SerialRequest::Write { port, data, base64 } => {
        self.handle_write_port(sub_id, port, data, base64.unwrap_or(false))
      }
      SerialRequest::Close { port } => self.handle_close_port(sub_id, port),
      SerialRequest::List {} => self.handle_list_ports(sub_id),
    };
    if let Err(e) = response {
      warn!("Error '{}' occured handling serial request message", e);
      // Send error?
      self.send_message(&sub_id, to_serial_response_error(e));

    }
  }

  /// Handle write port requests
  fn handle_write_port(
    &mut self,
    sub_id: &String,
    port_name: String,
    data: String,
    base_64: bool,
  ) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    self.check_owns_writelock(&port_name, &sub_id)?;
    match base_64 {
      true => {
        base64::decode(&data)
          .map_err(|e| ErrorKind::Base64(e).into())
          .and_then(|d| self.port_manager.write_port(&port_name, &d))
          .map(|_| self.send_message(&sub_id, SerialResponse::Wrote { port: port_name }),)
      }
      false => {
        self
          .port_manager
          .write_port(&port_name, data.as_bytes())
          .map(|_| self.send_message(&sub_id, SerialResponse::Wrote { port: port_name }),)
      }
    }
  }

  /// Handle write lock requests
  fn handle_write_lock(&mut self, sub_id: &String, port_name: String) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    self
      .writelock_manager
      .lock_port(&port_name, &sub_id)
      .map(|_| self.send_message(&sub_id, SerialResponse::WriteLocked { port: port_name }),)
  }

  /// Handle write requests
  fn handle_release_write_lock(
    &mut self,
    sub_id: &String,
    port_name: Option<String>,
  ) -> Result<()> {
    self.check_sub_id(sub_id)?;
    match port_name {
      None => {
        self.writelock_manager.unlock_all_ports_for_sub(&sub_id);
        Ok(
          self.send_message(
            &sub_id,
            SerialResponse::WriteLockReleased { port: port_name },
          ),
        )
      }
      Some(port_name) => {
        self
          .writelock_manager
          .unlock_port(&port_name, &sub_id)
          .map(
            |_| {
              self.send_message(
                &sub_id,
                SerialResponse::WriteLockReleased { port: Some(port_name) },
              )
            },
          )
      }
    }
  }

  /// Handle open port requests
  fn handle_open_port(&mut self, sub_id: &String, port_name: String) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    self.port_manager.open_port(&port_name)?;
    self
      .sub_manager
      .add_port(&sub_id, &port_name)
      .map(|_| self.send_message(&sub_id, SerialResponse::Opened { port: port_name }),)
  }

  /// Handle list ports request
  fn handle_list_ports(&mut self, sub_id: &String) -> Result<()> {
    self.check_sub_id(&sub_id)?;
    let port_names: Result<Vec<String>> =
      self
        .port_manager
        .list_ports()
        .map(|v| v.iter().map(|v| v.port_name.clone()).collect());
    let resp = port_names.map(|pns| SerialResponse::List { ports: pns });
    self.send_message(
      &sub_id,
      resp.unwrap_or(
        SerialResponse::Error {
          display: "".to_string(),
          description: "".to_string(),
        },
      ),
    );
    Ok(())
  }

  /// Handle close port requests
  fn handle_close_port(&mut self, sub_id: &String, port_name: Option<String>) -> Result<()> {
    match port_name {
      Some(port_name) => self.handle_close_port_for_sub(sub_id, port_name),
      None => self.handle_close_all_ports_for_sub(sub_id)
    }
  }

  /// Handle closing a signle port for a sub
  fn handle_close_port_for_sub(&mut self, sub_id: &String, port_name: String) -> Result<()> {
    self.sub_manager.remove_port(&sub_id, &port_name)?;
    self.writelock_manager.unlock_port_if_locked_by(&port_name, &sub_id);
    // self.cleanup_ports_with_no_subs();
    let close_resp = SerialResponse::Closed { port: port_name.clone() };
    self.send_message(&sub_id, close_resp);
    Ok(())
  }

  /// Handle closing all ports for sub
  fn handle_close_all_ports_for_sub(&mut self, sub_id: &String) -> Result<()> {
    self.sub_manager.clear_ports(Some(&sub_id));
    self.writelock_manager.unlock_all_ports_for_sub(sub_id);

    // Close ports with no subscribers
    let open_ports = self.port_manager.open_ports();
    let subscribed_ports = self.sub_manager.subscribed_ports();
    let ports_with_no_subs = open_ports.difference(&subscribed_ports);

    // For each open port that isn't subscribed,
    for port_to_close in ports_with_no_subs {
      // close it, REDUNDANT?
      self.port_manager.close_port(&port_to_close);
      // remove the write lock, REDUNDANT?
      self.writelock_manager.clear_lock(&port_to_close);
      // Let them know its closed
      let close_resp = SerialResponse::Closed { port: port_to_close.clone() };
      self.send_message(&sub_id, close_resp);
    }
    Ok(())
  }

  /// Cleanup any bad ports
  fn cleanup_bad_ports(&mut self, bad_ports: &HashSet<String>) {
    for port_name in bad_ports.iter() {
      // Tell everyone port is sick
      let err_resp = to_serial_response_error(ErrorKind::PortReadError(port_name.to_owned()).into(),);
      self.broadcast_message_for_port(port_name, err_resp);
      // Tell everyone the sick ports were closed
      let close_resp = SerialResponse::Closed { port: port_name.clone() };
      self.broadcast_message_for_port(port_name, close_resp);
      // Close bad ports
      self.port_manager.close_port(port_name);
      // Remove write locks on bad ports
      self.writelock_manager.clear_lock(port_name);
      // Remove bad ports from subscriptions
      self.sub_manager.remove_port_from_all(port_name);
    }
  }

  fn send_message(&mut self, sub_id: &String, msg: SerialResponse) {
    if let Err(e) = self.sub_manager.send_message(sub_id, msg) {
      warn!("Error sending serial response to sub_id '{}'", sub_id);
      let mut bad_subs = Vec::new();
      bad_subs.push(e);
      self.cleanup_bad_subs(bad_subs);
    }
  }

  /// Broadcast a message to all subscribers and
  /// then cleanup any subs that errored
  fn broadcast_message(&mut self, msg: SerialResponse) {
    let bad_subs = self.sub_manager.broadcast_message(msg);
    self.cleanup_bad_subs(bad_subs);
  }

  /// Broadcast a message to all subscribers interested in the
  /// given port and then cleanup any subs that errored
  fn broadcast_message_for_port(&mut self, port_name: &String, msg: SerialResponse) {
    let bad_subs = self
      .sub_manager
      .broadcast_message_for_port(port_name, msg);
    self.cleanup_bad_subs(bad_subs);
  }

  /// Cleanup any bad subs where a message send failed
  fn cleanup_bad_subs(&mut self, bad_subs: Vec<Error>) {
    for e in bad_subs {
      if let Error(ErrorKind::SubscriberSendError(sub_id), _) = e {
        // Remove subscriptions
        self.sub_manager.end_subscription(&sub_id);
        // Remove all write locks held by dead subscription
        self.writelock_manager.unlock_all_ports_for_sub(&sub_id);
      }
    }
  }

  fn check_sub_id(&self, sub_id: &String) -> Result<()> {
    self.sub_manager.check_subscription_exists(sub_id)
  }

  fn check_owns_writelock(&self, port_name: &String, sub_id: &String) -> Result<()> {
    self
      .writelock_manager
      .check_owns_write_lock(port_name, sub_id)
  }
}
<|endoftext|>


<commit_before>

use std::cell::RefCell;
use std::collections::HashMap;
use std::rc::Rc;
use std::sync::mpsc::{Sender, Receiver, TryRecvError};
use std::thread;
use std::time::Duration;

use serialport as sp;
use base64;

use serial_support::messages::*;

/// Convenience type for a listener
/// that accepts weak refs of Senders of Serial Reponses
/// This is how the manager will communicate
/// results back to the websockets
type SubscReceiver = Receiver<SubscriptionRequest>;

/// Struct for containing Port information
pub struct OpenPort {
  /// Handle that controls writes to port
  write_lock_sub_id: String,
  /// The opened serial port
  /// SerialPort is not Sized, so it makes hashmap mad
  /// and so we deal with these shennanigans
  port: Rc<RefCell<sp::SerialPort>>,
}

pub struct Subscription {
  /// Subscription
  subscriber: Sender<SerialResponse>,
  /// The ports it is subscribed to
  ports: Vec<String>,
  /// Subscription id
  sub_id: String,
}

/// Manages tracking and reading/writing from serial
/// ports
struct SerialPortManager {
  /// List of port names to subscribers
  subscriptions: Vec<Subscription>,
  /// Receiver for serial requests
  receiver: Receiver<SerialRequest>,
  /// Receiver for response subscription requests
  subsc_receiver: SubscReceiver,
  /// Maintains list of ports
  open_ports: HashMap<String, OpenPort>,
}

impl SerialPortManager {
  fn open_port(&self, sub_id: String, port: String) {
    // Check if port is already open

    // If not, open it
    let sp_settings = sp::SerialPortSettings {
      baud_rate: sp::BaudRate::Baud115200,
      data_bits: sp::DataBits::Eight,
      flow_control: sp::FlowControl::None,
      parity: sp::Parity::None,
      stop_bits: sp::StopBits::One,
      timeout: Duration::from_millis(1),
    };

  }

  fn add_sender(&self, port: String, sender: Sender<SerialResponse>) {

  }

  fn create_write_lock(&self, sub_id: String, port: String) {

  }

  fn release_write_lock(&self, sub_id: String, port: Option<String>) {

  }

  fn write_port(&self, sub_id: String, port: String, data: String, base64: bool) {

  }

  fn close_port(&self, sub_id: String, port: Option<String>) {}

  /// Handles and dispatches SerialRequest sent by
  /// the channel
  fn handle_serial_request(&self, msg: SerialRequest) {
    match msg {
      SerialRequest::Open { sub_id, port } => self.open_port(sub_id, port),
      SerialRequest::WriteLock { sub_id, port } => self.add_write_lock(sub_id, port),
      SerialRequest::ReleaseWriteLock { sub_id, port } => self.release_write_lock(sub_id, port),
      SerialRequest::Write {
        sub_id,
        port,
        data,
        base64,
      } => self.write_port(sub_id, port, data, base64.unwrap_or(false)),
      SerialRequest::Close { sub_id, port } => self.close_port(sub_id, port),
    }
  }

  /// Fire up the port manager
  ///
  /// TODO Right now this is spinning a thread between reading
  /// subscription requests, reading serial port data
  /// etc. This will cause high cpu usage
  ///
  /// TODO Once this is working, refactor to avoid
  /// thread spinning?
  ///
  fn run(&mut self) {

    let mut shutdown = false;
    // Check about 30 times a second
    let sleep_dur = Duration::from_millis(33);
    let mut serial_buf: Vec<u8> = vec![0; 4096];

    while !shutdown {
      // Sleep for a little bit to avoid pegging cpu
      thread::sleep(sleep_dur);

      // Handle serial operation requests
      match self.receiver.try_recv() {
        Err(e) => {
          match e {
            TryRecvError::Empty => {
              // nothing to do
            }
            TryRecvError::Disconnected => {
              // Remote end hung up, time to shutdown
              info!("Shutting down SerialPortManager");
            }
          }
        }
        Ok(req) => self.handle_serial_request(req),
      }

      // Check for new data on each port
      for (port_name, open_port) in self.open_ports.iter() {
        // If we have data, send it back
        // Read needs to borrow &mut self, so wrap port in RC?
        let result = match open_port
                .port
                .borrow_mut()
                .read(serial_buf.as_mut_slice()) {
          Ok(bytes_read) => {
            // We got some data
            let bytes = &serial_buf[0..bytes_read];
            // Try and parse the bytes as utf-8
            match String::from_utf8(bytes.to_vec()) {
              // We need to send as binary
              Err(_) => {
                SerialResponse::Read {
                  port: port_name.to_string(),
                  data: base64::encode(bytes),
                  base64: Some(true),
                }
              }
              Ok(s) => {
                SerialResponse::Read {
                  port: port_name.to_string(),
                  data: s,
                  base64: Some(false),
                }
              }
            }
          }
          Err(_) => {
            // We failed to read the port, send error
            SerialResponse::Error {
              kind: ErrorType::ReadError,
              msg: "Error reading serial port '{}'".to_string(),
              detail: None,
              port: Some(port_name.to_string()),
              sub_id: None,
            }
          }
        };

        // Send report, retain only live connections
        self
          .subscriptions
          .retain(|c| {
            // If its not the port we are currently interested in, we keep it
            if !c.ports.contains(&port_name) {
              true
            } else {
              // Otherwise try and send on it, and if it fails
              // remove the subscription
              c.subscriber
                .send(result.clone())
                .map(|_| true)
                .map_err(|_| {
                           warn!("Connection closed.");
                           false
                         })
                .unwrap()
            }
          });
      }

      // Check for Subscribe requests
      match self.subsc_receiver.try_recv() {
        Err(e) => {
          match e {
            TryRecvError::Empty => {
              // nothing to do
            }
            TryRecvError::Disconnected => {
              // Remote end hung up, time to shutdown
              shutdown = true;
              info!("Shutting down SerialPortManager");
            }
          }
        }
        Ok(subReq) => {
          self
            .subscriptions
            .push(Subscription {
                    ports: Vec::with_capacity(4),
                    subscriber: subReq.subscriber,
                    sub_id: subReq.sub_id,
                  });
        }
      }
    }
  }
}

// TODO How to spawn and fire it up
// TODO Document that closing subscription sender will result in
// portmanager shutdown.
// fn spawnManager() -> {}
<commit_msg>Adding convenience methods to avoid spaghetti code<commit_after>

use std::cell::RefCell;
use std::collections::HashMap;
use std::rc::Rc;
use std::sync::mpsc::{Sender, Receiver, TryRecvError};
use std::thread;
use std::time::Duration;

use serialport as sp;
use base64;

use serial_support::messages::*;

/// Convenience type for a listener
/// that accepts weak refs of Senders of Serial Reponses
/// This is how the manager will communicate
/// results back to the websockets
type SubscReceiver = Receiver<SubscriptionRequest>;

/// Struct for containing Port information
pub struct OpenPort {
  /// Handle that controls writes to port
  write_lock_sub_id: Option<String>,
  /// The opened serial port
  /// SerialPort is not Sized, so it makes hashmap mad
  /// and so we deal with these shennanigans
  port: RefCell<Box<sp::SerialPort>>,
}

/// Subscription
pub struct Subscription {
  /// Subscription
  subscriber: Sender<SerialResponse>,
  /// The ports it is subscribed to
  ports: Vec<String>,
  /// Subscription id
  sub_id: String,
}

/// Manages tracking and reading/writing from serial
/// ports
///
/// TODO efficiency can be improved through choice of better
/// data structures though looping through a short list
/// is probably faster than hashing...
struct SerialPortManager {
  /// Receiver for serial requests
  receiver: Receiver<SerialRequest>,
  /// Receiver for response subscription requests
  subsc_receiver: SubscReceiver,
  /// List of port names to subscribers
  subscriptions: Vec<Subscription>,
  /// Maintains list of ports
  open_ports: HashMap<String, OpenPort>,
}

impl SerialPortManager {

  /// Try and get the subscription for sub_id
  /// If unsuccessful, return a Err<SerialResponse::Error>
  fn get_subscription(&mut self, sub_id: String) -> Result<&Subscription, SerialResponse> {
    match self.subscriptions.iter().find(|&s| s.sub_id == sub_id ){
      None => Err(SerialResponse::Error{
        kind: ErrorType::PortNotFound,
        msg: format!("Could not find subscription for '{}'", sub_id),
        detail: None,
        port: None,
        sub_id: Some(sub_id)
      }),
      Some(sub) => Ok(sub)
    }
  }

  /// Try and get the serial port
  /// If unsuccessful, return a Err<SerialResponse::Error>
  fn get_port(&mut self, port:String, sub_id: Option<String>) -> Result<&OpenPort,SerialResponse>{
    match self.open_ports.get(&port){
      None => Err(SerialResponse::Error{
        kind: ErrorType::PortNotFound,
        msg: format!("Could not find open serial port for '{}'", port),
        detail: None,
        port: Some(port),
        sub_id: sub_id
      }),
      Some(sp) => Ok(sp)
    }
  }

  /// Open the given port.
  /// If the port is already open, then just subscribe to it
  fn open_port(&mut self, sub_id: String, port: String) {

    // If not, open it
    let sp_settings = sp::SerialPortSettings {
      baud_rate: sp::BaudRate::Baud115200,
      data_bits: sp::DataBits::Eight,
      flow_control: sp::FlowControl::None,
      parity: sp::Parity::None,
      stop_bits: sp::StopBits::One,
      timeout: Duration::from_millis(1),
    };

    // match self.subscriptions.iter().position(|s| s.sub_id == sub_id){
    //   // We don't even know who to send an error to...
    //   None => warn!("No subscription found for sub_id {}",sub_id),
    //   Some(idx) => {
    //     let sub = &mut self.subscriptions[idx];
    //     let serial_port = match self.open_ports.get(&port){
    //       None => {
    //         let sp = sp::open_with_settings(&port, &sp_settings);
    //         self.open_ports.insert(port, sp);
    //         return
    //       }
    //       Some(sp) => Ok(sp.port.borrow_mut())
    //     };

    //   }
    // }

    // match self.open_ports.get(&port) {
    //   case Some(serial_port) => {

    //   },
    //   case None => {

    //   }
    // }

  }

  // fn add_sender(&self, port: String, sender: Sender<SerialResponse>) {}

  fn set_write_lock(&self, sub_id: String, port: String) {

  }

  fn release_write_lock(&self, sub_id: String, port: Option<String>) {}

  fn write_port(&mut self, sub_id: String, port: String, data: String, base64: bool) {
    // get port

    // check sub_id matches write_lock id

    // If it does, try and write

    // Else send a write lock error
  }

  /// Send a response to a subscription with the id sub_id
  fn send_response(&mut self, sub_id: String, response: SerialResponse) {
    match self.subscriptions.iter().position(|s| s.sub_id == sub_id){
      Some(idx) => {
        match self.subscriptions[idx].subscriber.send(response.clone()){
          Ok(_) => debug!("Sucessfully sent {:?} to {}",response.clone(), sub_id),
          Err(_) => {
            debug!("Remote end for sub_id {} closed, removing", sub_id);
            self.subscriptions.remove(idx);
          }
        }
      }
      None => warn!("No subscription found for sub_id '{}'",sub_id)
    }
  }

  /// Broadcast a response to all subscriptions
  fn broadcast(&mut self, response: SerialResponse) {
    self
      .subscriptions
      .retain(|sub| {
        sub
          .subscriber
          .send(response.clone())
          .map(|_| true)
          .unwrap_or_else(|_| {
                     warn!("Connection closed.");
                     false
                   })
      });
  }

  fn close_port(&self, sub_id: String, port: Option<String>) {}

  /// Handles and dispatches SerialRequest sent by
  /// the channel
  fn handle_serial_request(&mut self, msg: SerialRequest) {
    match msg {
      SerialRequest::Open { sub_id, port } => self.open_port(sub_id, port),
      SerialRequest::WriteLock { sub_id, port } => self.set_write_lock(sub_id, port),
      SerialRequest::ReleaseWriteLock { sub_id, port } => self.release_write_lock(sub_id, port),
      SerialRequest::Write {
        sub_id,
        port,
        data,
        base64,
      } => self.write_port(sub_id, port, data, base64.unwrap_or(false)),
      SerialRequest::Close { sub_id, port } => self.close_port(sub_id, port),
    }
  }

  /// Fire up the port manager
  ///
  /// TODO Right now this is spinning a thread between reading
  /// subscription requests, reading serial port data
  /// etc. This will cause high cpu usage
  ///
  /// TODO Once this is working, refactor to avoid
  /// thread spinning?
  ///
  fn run(&mut self) {

    let mut shutdown = false;
    // Check about 30 times a second
    let sleep_dur = Duration::from_millis(33);
    let mut serial_buf: Vec<u8> = vec![0; 4096];

    while !shutdown {
      // Sleep for a little bit to avoid pegging cpu
      thread::sleep(sleep_dur);

      // Handle serial operation requests
      match self.receiver.try_recv() {
        Err(e) => {
          match e {
            TryRecvError::Empty => {
              // nothing to do
            }
            TryRecvError::Disconnected => {
              // Remote end hung up, time to shutdown
              info!("Shutting down SerialPortManager");
            }
          }
        }
        Ok(req) => self.handle_serial_request(req),
      }

      // Check for new data on each port
      for (port_name, open_port) in self.open_ports.iter() {
        // If we have data, send it back
        // Read needs to borrow &mut self, so wrap port in RC?
        let result = match open_port
                .port
                .borrow_mut()
                .read(serial_buf.as_mut_slice()) {
          Ok(bytes_read) => {
            // We got some data
            let bytes = &serial_buf[0..bytes_read];
            // Try and parse the bytes as utf-8
            match String::from_utf8(bytes.to_vec()) {
              // We need to send as binary
              Err(_) => {
                SerialResponse::Read {
                  port: port_name.to_string(),
                  data: base64::encode(bytes),
                  base64: Some(true),
                }
              }
              Ok(s) => {
                SerialResponse::Read {
                  port: port_name.to_string(),
                  data: s,
                  base64: Some(false),
                }
              }
            }
          }
          Err(_) => {
            // We failed to read the port, send error
            SerialResponse::Error {
              kind: ErrorType::ReadError,
              msg: "Error reading serial port '{}'".to_string(),
              detail: None,
              port: Some(port_name.to_string()),
              sub_id: None,
            }
          }
        };

        // Send report, retain only live connections
        self
          .subscriptions
          .retain(|c| {
            // If its not the port we are currently interested in, we keep it
            if !c.ports.contains(&port_name) {
              true
            } else {
              // Otherwise try and send on it, and if it fails
              // remove the subscription
              c.subscriber
                .send(result.clone())
                .map(|_| true)
                .unwrap_or_else(|_| {
                           warn!("Connection closed.");
                           false
                         })
            }
          });
      }

      // Check for Subscribe requests
      match self.subsc_receiver.try_recv() {
        Err(e) => {
          match e {
            TryRecvError::Empty => {
              // nothing to do
            }
            TryRecvError::Disconnected => {
              // Remote end hung up, time to shutdown
              shutdown = true;
              info!("Shutting down SerialPortManager");
            }
          }
        }
        Ok(sub_req) => {
          self
            .subscriptions
            .push(Subscription {
                    ports: Vec::with_capacity(4),
                    subscriber: sub_req.subscriber,
                    sub_id: sub_req.sub_id,
                  });
        }
      }
    }
  }
}

// TODO How to spawn and fire it up
// TODO Document that closing subscription sender will result in
// portmanager shutdown.
// fn spawnManager() -> {}
<|endoftext|>


<commit_before>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

#![allow(dead_code)]

use server::Node;
use server::transport::{ServerRaftStoreRouter, RaftStoreRouter};
use raftstore::store::{Transport, Peekable};
use raftstore::errors::Error as RaftServerError;
use raftstore::coprocessor::RegionSnapshot;
use util::HandyRwLock;
use kvproto::raft_cmdpb::{RaftCmdRequest, RaftCmdResponse, RaftRequestHeader, Request, Response,
                          GetRequest, CmdType, SeekRequest, DeleteRequest, PutRequest};
use kvproto::errorpb;
use kvproto::kvrpcpb::Context;

use pd::PdClient;
use uuid::Uuid;
use std::sync::{Arc, RwLock};
use std::fmt::{self, Formatter, Debug};
use std::io::Error as IoError;
use std::time::Duration;
use std::result;
use rocksdb::DB;
use protobuf::RepeatedField;

use storage::engine;
use super::{Engine, Modify, Snapshot};
use util::event::Event;
use storage::{Key, Value, KvPair};

const DEFAULT_TIMEOUT_SECS: u64 = 5;

quick_error! {
    #[derive(Debug)]
    pub enum Error {
        RequestFailed(e: errorpb::Error) {
            from()
            description(e.get_message())
        }
        Io(e: IoError) {
            from()
            cause(e)
            description(e.description())
        }
        RocksDb(reason: String) {
            description(reason)
        }
        Server(e: RaftServerError) {
            from()
            cause(e)
            description(e.description())
        }
        InvalidResponse(reason: String) {
            description(reason)
        }
        InvalidRequest(reason: String) {
            description(reason)
        }
        Timeout(d: Duration) {
            description("request timeout")
            display("timeout after {:?}", d)
        }
    }
}

pub type Result<T> = result::Result<T, Error>;

impl From<Error> for engine::Error {
    fn from(e: Error) -> engine::Error {
        match e {
            Error::RequestFailed(e) => engine::Error::Request(e),
            e => box_err!(e),
        }
    }
}

/// RaftKv is a storage engine base on RaftKvServer.
pub struct RaftKv<T: PdClient + 'static, Trans: Transport + 'static> {
    node: Node<T, Trans>,
    db: Arc<DB>,
    router: Arc<RwLock<ServerRaftStoreRouter>>,
}

impl<T: PdClient, Trans: Transport> RaftKv<T, Trans> {
    /// Create a RaftKv using specified configuration.
    pub fn new(node: Node<T, Trans>, db: Arc<DB>) -> RaftKv<T, Trans> {
        let router = node.raft_store_router();
        RaftKv {
            node: node,
            db: db,
            router: router,
        }
    }

    pub fn call_command(&self, request: RaftCmdRequest) -> Result<Event<RaftCmdResponse>> {
        let finished = Event::new();
        let finished2 = finished.clone();
        let timeout = Duration::from_secs(DEFAULT_TIMEOUT_SECS);

        try!(self.router.rl().send_command(request,
                                           box move |resp| {
                                               finished2.set(resp);
                                               // Wait for response to be consumed or `finished` is
                                               // dropped.
                                               finished2.wait_clear(None);
                                               Ok(())
                                           }));

        if finished.wait_timeout(Some(timeout)) {
            return Ok(finished);
        }

        Err(Error::Timeout(timeout))
    }

    fn exec_cmd_request(&self, req: RaftCmdRequest) -> Result<Event<RaftCmdResponse>> {
        let uuid = req.get_header().get_uuid().to_vec();
        let l = req.get_requests().len();

        let resp = try!(self.call_command(req));
        try!(resp.apply(|resp| {
                     if resp.get_header().get_uuid() != &*uuid {
                         return Err(Error::InvalidResponse("response is not correct!!!"
                                                               .to_owned()));
                     }
                     if resp.get_header().has_error() {
                         return Err(Error::RequestFailed(resp.take_header().take_error()));
                     }
                     if l != resp.get_responses().len() {
                         return Err(Error::InvalidResponse("response count is not equal to \
                                                            requests, something must go wrong."
                                                               .to_owned()));
                     }
                     Ok(())
                 })
                 .unwrap());
        Ok(resp)
    }

    fn new_request_header(&self, ctx: &Context) -> RaftRequestHeader {
        let mut header = RaftRequestHeader::new();
        header.set_region_id(ctx.get_region_id());
        header.set_region_epoch(ctx.get_region_epoch().clone());
        header.set_uuid(Uuid::new_v4().as_bytes().to_vec());
        header
    }

    fn exec_requests(&self, ctx: &Context, reqs: Vec<Request>) -> Result<Vec<Response>> {
        let header = self.new_request_header(ctx);
        let mut cmd = RaftCmdRequest::new();
        cmd.set_header(header);
        cmd.set_requests(RepeatedField::from_vec(reqs));
        let resp = try!(self.exec_cmd_request(cmd));
        Ok(resp.take().unwrap().take_responses().to_vec())
    }

    fn exec_request(&self, ctx: &Context, req: Request) -> Result<Response> {
        let resps = self.exec_requests(ctx, vec![req]);
        resps.map(|mut v| v.remove(0))
    }
}

fn invalid_resp_type(exp: CmdType, act: CmdType) -> engine::Error {
    Error::InvalidResponse(format!("cmd type not match, want {:?}, got {:?}!", exp, act)).into()
}

impl<T: PdClient, Trans: Transport> Debug for RaftKv<T, Trans> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "RaftKv")
    }
}

impl<T: PdClient, Trans: Transport> Engine for RaftKv<T, Trans> {
    fn get(&self, ctx: &Context, key: &Key) -> engine::Result<Option<Value>> {
        let mut get = GetRequest::new();
        get.set_key(key.raw().clone());
        let mut req = Request::new();
        req.set_cmd_type(CmdType::Get);
        req.set_get(get);
        let mut resp = try!(self.exec_request(ctx, req));
        if resp.get_cmd_type() != CmdType::Get {
            return Err(invalid_resp_type(CmdType::Get, resp.get_cmd_type()));
        }
        let get_resp = resp.mut_get();
        if get_resp.has_value() {
            Ok(Some(get_resp.take_value()))
        } else {
            Ok(None)
        }
    }

    fn seek(&self, ctx: &Context, key: &Key) -> engine::Result<Option<KvPair>> {
        let mut seek = SeekRequest::new();
        seek.set_key(key.raw().clone());
        let mut req = Request::new();
        req.set_cmd_type(CmdType::Seek);
        req.set_seek(seek);
        let mut resp = try!(self.exec_request(ctx, req));
        if resp.get_cmd_type() != CmdType::Seek {
            return Err(invalid_resp_type(CmdType::Seek, resp.get_cmd_type()));
        }
        let seek_resp = resp.mut_seek();
        if seek_resp.has_key() {
            Ok(Some((seek_resp.take_key(), seek_resp.take_value())))
        } else {
            Ok(None)
        }
    }

    fn write(&self, ctx: &Context, mut modifies: Vec<Modify>) -> engine::Result<()> {
        if modifies.len() == 0 {
            return Ok(());
        }
        let mut reqs = Vec::with_capacity(modifies.len());
        while !modifies.is_empty() {
            let m = modifies.pop().unwrap();
            let mut req = Request::new();
            match m {
                Modify::Delete(k) => {
                    let mut delete = DeleteRequest::new();
                    delete.set_key(k.raw().clone());
                    req.set_cmd_type(CmdType::Delete);
                    req.set_delete(delete);
                }
                Modify::Put((k, v)) => {
                    let mut put = PutRequest::new();
                    put.set_key(k.raw().clone());
                    put.set_value(v);
                    req.set_cmd_type(CmdType::Put);
                    req.set_put(put);
                }
            }
            reqs.push(req);
        }
        try!(self.exec_requests(ctx, reqs));
        Ok(())
    }

    fn snapshot<'a>(&'a self, ctx: &Context) -> engine::Result<Box<Snapshot + 'a>> {
        let mut req = Request::new();
        req.set_cmd_type(CmdType::Snap);

        let header = self.new_request_header(ctx);
        let mut cmd = RaftCmdRequest::new();
        cmd.set_header(header);
        cmd.mut_requests().push(req);

        let resp = try!(self.exec_cmd_request(cmd));
        let region = try!(resp.apply(|resp| {
                                  let mut resp = resp.take_responses().remove(0);
                                  if resp.get_cmd_type() != CmdType::Snap {
                                      return Err(invalid_resp_type(CmdType::Snap,
                                                                   resp.get_cmd_type()));
                                  }
                                  Ok(resp.take_snap().take_region().clone())
                              })
                              .unwrap());
        let snap = RegionSnapshot::from_raw(self.db.as_ref(), region);
        Ok(box snap)
    }
}

impl<'a> Snapshot for RegionSnapshot<'a> {
    fn get(&self, key: &Key) -> engine::Result<Option<Value>> {
        let v = box_try!(self.get_value(key.raw()));
        Ok(v.map(|v| v.to_vec()))
    }

    fn seek(&self, key: &Key) -> engine::Result<Option<KvPair>> {
        let pair = box_try!(self.seek(key.raw()));
        Ok(pair)
    }

    fn reverse_seek(&self, key: &Key) -> engine::Result<Option<KvPair>> {
        let pair = box_try!(self.reverse_seek(key.raw()));
        Ok(pair)
    }
}
<commit_msg>raftkv: add comments. (#545)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

#![allow(dead_code)]

use server::Node;
use server::transport::{ServerRaftStoreRouter, RaftStoreRouter};
use raftstore::store::{Transport, Peekable};
use raftstore::errors::Error as RaftServerError;
use raftstore::coprocessor::RegionSnapshot;
use util::HandyRwLock;
use kvproto::raft_cmdpb::{RaftCmdRequest, RaftCmdResponse, RaftRequestHeader, Request, Response,
                          GetRequest, CmdType, SeekRequest, DeleteRequest, PutRequest};
use kvproto::errorpb;
use kvproto::kvrpcpb::Context;

use pd::PdClient;
use uuid::Uuid;
use std::sync::{Arc, RwLock};
use std::fmt::{self, Formatter, Debug};
use std::io::Error as IoError;
use std::time::Duration;
use std::result;
use rocksdb::DB;
use protobuf::RepeatedField;

use storage::engine;
use super::{Engine, Modify, Snapshot};
use util::event::Event;
use storage::{Key, Value, KvPair};

const DEFAULT_TIMEOUT_SECS: u64 = 5;

quick_error! {
    #[derive(Debug)]
    pub enum Error {
        RequestFailed(e: errorpb::Error) {
            from()
            description(e.get_message())
        }
        Io(e: IoError) {
            from()
            cause(e)
            description(e.description())
        }
        RocksDb(reason: String) {
            description(reason)
        }
        Server(e: RaftServerError) {
            from()
            cause(e)
            description(e.description())
        }
        InvalidResponse(reason: String) {
            description(reason)
        }
        InvalidRequest(reason: String) {
            description(reason)
        }
        Timeout(d: Duration) {
            description("request timeout")
            display("timeout after {:?}", d)
        }
    }
}

pub type Result<T> = result::Result<T, Error>;

impl From<Error> for engine::Error {
    fn from(e: Error) -> engine::Error {
        match e {
            Error::RequestFailed(e) => engine::Error::Request(e),
            e => box_err!(e),
        }
    }
}

/// RaftKv is a storage engine base on RaftKvServer.
pub struct RaftKv<T: PdClient + 'static, Trans: Transport + 'static> {
    node: Node<T, Trans>,
    db: Arc<DB>,
    router: Arc<RwLock<ServerRaftStoreRouter>>,
}

impl<T: PdClient, Trans: Transport> RaftKv<T, Trans> {
    /// Create a RaftKv using specified configuration.
    pub fn new(node: Node<T, Trans>, db: Arc<DB>) -> RaftKv<T, Trans> {
        let router = node.raft_store_router();
        RaftKv {
            node: node,
            db: db,
            router: router,
        }
    }

    pub fn call_command(&self, request: RaftCmdRequest) -> Result<Event<RaftCmdResponse>> {
        let finished = Event::new();
        let finished2 = finished.clone();
        let timeout = Duration::from_secs(DEFAULT_TIMEOUT_SECS);

        try!(self.router.rl().send_command(request,
                                           box move |resp| {
                                               finished2.set(resp);
                                               // Wait for response to be consumed or `finished` is
                                               // dropped.
                                               finished2.wait_clear(None);
                                               Ok(())
                                           }));

        if finished.wait_timeout(Some(timeout)) {
            return Ok(finished);
        }

        Err(Error::Timeout(timeout))
    }

    fn exec_cmd_request(&self, req: RaftCmdRequest) -> Result<Event<RaftCmdResponse>> {
        let uuid = req.get_header().get_uuid().to_vec();
        let l = req.get_requests().len();

        let resp = try!(self.call_command(req));
        try!(resp.apply(|resp| {
                     if resp.get_header().get_uuid() != &*uuid {
                         return Err(Error::InvalidResponse("response is not correct!!!"
                                                               .to_owned()));
                     }
                     if resp.get_header().has_error() {
                         return Err(Error::RequestFailed(resp.take_header().take_error()));
                     }
                     if l != resp.get_responses().len() {
                         return Err(Error::InvalidResponse("response count is not equal to \
                                                            requests, something must go wrong."
                                                               .to_owned()));
                     }
                     Ok(())
                 })
                 .unwrap());
        Ok(resp)
    }

    fn new_request_header(&self, ctx: &Context) -> RaftRequestHeader {
        let mut header = RaftRequestHeader::new();
        header.set_region_id(ctx.get_region_id());
        header.set_region_epoch(ctx.get_region_epoch().clone());
        header.set_uuid(Uuid::new_v4().as_bytes().to_vec());
        header
    }

    fn exec_requests(&self, ctx: &Context, reqs: Vec<Request>) -> Result<Vec<Response>> {
        let header = self.new_request_header(ctx);
        let mut cmd = RaftCmdRequest::new();
        cmd.set_header(header);
        cmd.set_requests(RepeatedField::from_vec(reqs));
        let resp = try!(self.exec_cmd_request(cmd));
        // notice: raft thread will be woken up after the take method is called.
        Ok(resp.take().unwrap().take_responses().to_vec())
    }

    fn exec_request(&self, ctx: &Context, req: Request) -> Result<Response> {
        let resps = self.exec_requests(ctx, vec![req]);
        resps.map(|mut v| v.remove(0))
    }
}

fn invalid_resp_type(exp: CmdType, act: CmdType) -> engine::Error {
    Error::InvalidResponse(format!("cmd type not match, want {:?}, got {:?}!", exp, act)).into()
}

impl<T: PdClient, Trans: Transport> Debug for RaftKv<T, Trans> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "RaftKv")
    }
}

impl<T: PdClient, Trans: Transport> Engine for RaftKv<T, Trans> {
    fn get(&self, ctx: &Context, key: &Key) -> engine::Result<Option<Value>> {
        let mut get = GetRequest::new();
        get.set_key(key.raw().clone());
        let mut req = Request::new();
        req.set_cmd_type(CmdType::Get);
        req.set_get(get);
        let mut resp = try!(self.exec_request(ctx, req));
        if resp.get_cmd_type() != CmdType::Get {
            return Err(invalid_resp_type(CmdType::Get, resp.get_cmd_type()));
        }
        let get_resp = resp.mut_get();
        if get_resp.has_value() {
            Ok(Some(get_resp.take_value()))
        } else {
            Ok(None)
        }
    }

    fn seek(&self, ctx: &Context, key: &Key) -> engine::Result<Option<KvPair>> {
        let mut seek = SeekRequest::new();
        seek.set_key(key.raw().clone());
        let mut req = Request::new();
        req.set_cmd_type(CmdType::Seek);
        req.set_seek(seek);
        let mut resp = try!(self.exec_request(ctx, req));
        if resp.get_cmd_type() != CmdType::Seek {
            return Err(invalid_resp_type(CmdType::Seek, resp.get_cmd_type()));
        }
        let seek_resp = resp.mut_seek();
        if seek_resp.has_key() {
            Ok(Some((seek_resp.take_key(), seek_resp.take_value())))
        } else {
            Ok(None)
        }
    }

    fn write(&self, ctx: &Context, mut modifies: Vec<Modify>) -> engine::Result<()> {
        if modifies.len() == 0 {
            return Ok(());
        }
        let mut reqs = Vec::with_capacity(modifies.len());
        while !modifies.is_empty() {
            let m = modifies.pop().unwrap();
            let mut req = Request::new();
            match m {
                Modify::Delete(k) => {
                    let mut delete = DeleteRequest::new();
                    delete.set_key(k.raw().clone());
                    req.set_cmd_type(CmdType::Delete);
                    req.set_delete(delete);
                }
                Modify::Put((k, v)) => {
                    let mut put = PutRequest::new();
                    put.set_key(k.raw().clone());
                    put.set_value(v);
                    req.set_cmd_type(CmdType::Put);
                    req.set_put(put);
                }
            }
            reqs.push(req);
        }
        try!(self.exec_requests(ctx, reqs));
        Ok(())
    }

    fn snapshot<'a>(&'a self, ctx: &Context) -> engine::Result<Box<Snapshot + 'a>> {
        let mut req = Request::new();
        req.set_cmd_type(CmdType::Snap);

        let header = self.new_request_header(ctx);
        let mut cmd = RaftCmdRequest::new();
        cmd.set_header(header);
        cmd.mut_requests().push(req);

        // notice: the raft thread will be woken up after resp is dropped
        let resp = try!(self.exec_cmd_request(cmd));
        let region = try!(resp.apply(|resp| {
                                  let mut resp = resp.take_responses().remove(0);
                                  if resp.get_cmd_type() != CmdType::Snap {
                                      return Err(invalid_resp_type(CmdType::Snap,
                                                                   resp.get_cmd_type()));
                                  }
                                  Ok(resp.take_snap().take_region().clone())
                              })
                              .unwrap());

        let snap = RegionSnapshot::from_raw(self.db.as_ref(), region);
        Ok(box snap)
    }
}

impl<'a> Snapshot for RegionSnapshot<'a> {
    fn get(&self, key: &Key) -> engine::Result<Option<Value>> {
        let v = box_try!(self.get_value(key.raw()));
        Ok(v.map(|v| v.to_vec()))
    }

    fn seek(&self, key: &Key) -> engine::Result<Option<KvPair>> {
        let pair = box_try!(self.seek(key.raw()));
        Ok(pair)
    }

    fn reverse_seek(&self, key: &Key) -> engine::Result<Option<KvPair>> {
        let pair = box_try!(self.reverse_seek(key.raw()));
        Ok(pair)
    }
}
<|endoftext|>


<commit_before>use std::borrow::Cow;
use std::ffi::OsStr;
use std::path::{Path, PathBuf};

use globset;
use log::debug;
use walkdir;

use crate::template::{self, TemplateInput};
use crate::version::{self, Version};

use super::{config, VutConfig, VutError};

pub fn generate_template_output(
    config: &VutConfig,
    root_path: &Path,
    version: &Version,
    dir_entries: &[walkdir::DirEntry],
) -> Result<(), VutError> {
    let template_globsets = build_template_globsets(config)?;

    let template_input = generate_template_input(version)?;

    let mut processed_files: Vec<PathBuf> = Vec::new();
    let mut generated_files: Vec<PathBuf> = Vec::new();

    for (template, globset) in template_globsets {
        debug!("{:?}", template);

        let start_path: Cow<Path> = match &template.start_path {
            Some(p) => Cow::Owned(root_path.join(p)),
            None => Cow::Borrowed(root_path),
        };

        let output_path = match &template.output_path {
            Some(p) => Cow::Owned(root_path.join(p)),
            None => start_path.clone(),
        };

        let processor = template.processor.as_ref().map(|s| s.as_str()).unwrap_or("vut");
        let encoding = template.encoding.as_ref().map(|s| s.as_str());

        let template_files_iter = dir_entries
            .iter()
            .map(|entry| entry.path())
            // Exclude files outside the start path
            .filter(|path| path.starts_with(&start_path))
            // Only include template files
            .filter_map(|path| {
                // Make path relative, as we only want to match on the path
                // relative to the start path.
                let rel_path = path.strip_prefix(&start_path).unwrap();

                if globset.is_match(rel_path) {
                    Some((path, rel_path))
                } else {
                    None
                }
            });

        for (path, rel_path) in template_files_iter {
            debug!("Processing template {}", path.display());

            let output_file_name: &OsStr = path.file_stem().unwrap();
            let output_file_path = output_path.join(rel_path.with_file_name(output_file_name));

            template::generate_template_with_processor_name(
                processor,
                path,
                &output_file_path,
                &template_input,
                encoding,
            )
            .map_err(|err| VutError::TemplateGenerate(err))?;

            processed_files.push(path.to_path_buf());
            generated_files.push(output_file_path);
        }
    }

    Ok(())
}

pub fn generate_template_input(version: &Version) -> Result<TemplateInput, VutError> {
    let mut template_input = TemplateInput::new();

    let split_prerelease = version
        .prerelease
        .as_ref()
        .map_or(None, |p| version::split_numbered_prerelease(p));
    let split_build = version
        .build
        .as_ref()
        .map_or(None, |b| version::split_numbered_prerelease(b));

    template_input
        .values
        .insert("FullVersion".to_owned(), version.to_string());
    template_input.values.insert(
        "Version".to_owned(),
        Version {
            build: None,
            ..version.clone()
        }
        .to_string(),
    );
    template_input.values.insert(
        "MajorMinorPatch".to_owned(),
        format!("{}.{}.{}", version.major, version.minor, version.patch),
    );
    template_input
        .values
        .insert("MajorMinor".to_owned(), format!("{}.{}", version.major, version.minor));
    template_input
        .values
        .insert("Major".to_owned(), format!("{}", version.major));
    template_input
        .values
        .insert("Minor".to_owned(), format!("{}", version.minor));
    template_input
        .values
        .insert("Patch".to_owned(), format!("{}", version.patch));
    template_input.values.insert(
        "Prerelease".to_owned(),
        version.prerelease.as_ref().map_or("", |p| p).to_owned(),
    );
    template_input.values.insert(
        "PrereleasePrefix".to_owned(),
        split_prerelease
            .and_then(|sp| Some(sp.0.to_owned()))
            .unwrap_or_else(|| "".to_owned()),
    );
    template_input.values.insert(
        "PrereleaseNumber".to_owned(),
        split_prerelease
            .and_then(|sp| Some(format!("{}", sp.1)))
            .unwrap_or_else(|| "".to_owned()),
    );
    template_input
        .values
        .insert("Build".to_owned(), version.build.as_ref().map_or("", |b| b).to_owned());
    template_input.values.insert(
        "BuildPrefix".to_owned(),
        split_build
            .and_then(|sp| Some(sp.0.to_owned()))
            .unwrap_or_else(|| "".to_owned()),
    );
    template_input.values.insert(
        "BuildNumber".to_owned(),
        split_build
            .and_then(|sp| Some(format!("{}", sp.1)))
            .unwrap_or_else(|| "".to_owned()),
    );

    Ok(template_input)
}

pub fn build_template_globsets(config: &VutConfig) -> Result<Vec<(&config::Template, globset::GlobSet)>, VutError> {
    let mut template_globsets: Vec<(&config::Template, globset::GlobSet)> = Vec::new();

    for template in config.template.iter() {
        let patterns = match &template.pattern {
            config::Patterns::Single(v) => vec![v],
            config::Patterns::Multiple(v) => v.iter().collect(),
        };

        let mut globset = globset::GlobSetBuilder::new();

        for pattern in patterns {
            let glob = globset::Glob::new(&pattern).map_err(|err| VutError::Other(Cow::Owned(err.to_string())))?;
            globset.add(glob);
        }

        let globset = globset
            .build()
            .map_err(|err| VutError::Other(Cow::Owned(err.to_string())))?;

        template_globsets.push((&template, globset));
    }

    Ok(template_globsets)
}
<commit_msg>refactor: Clean up generate_template.rs<commit_after>use std::borrow::Cow;
use std::ffi::OsStr;
use std::path::{Path, PathBuf};

use globset;
use log::debug;
use walkdir;

use crate::template::{self, TemplateInput};
use crate::version::{self, Version};

use super::{config, VutConfig, VutError};

#[derive(Debug)]
struct TemplateSpec<'a> {
    include_globset: globset::GlobSet,
    def: &'a config::Template,
}

impl<'a> TemplateSpec<'a> {
    pub fn from_config_template(cfg_t: &'a config::Template) -> Result<Self, VutError> {
        let patterns = match &cfg_t.pattern {
            config::Patterns::Single(v) => vec![v],
            config::Patterns::Multiple(v) => v.iter().collect(),
        };

        let mut globset = globset::GlobSetBuilder::new();

        for pattern in patterns {
            let glob = globset::Glob::new(&pattern).map_err(|err| VutError::Other(Cow::Owned(err.to_string())))?;
            globset.add(glob);
        }

        let include_globset = globset
            .build()
            .map_err(|err| VutError::Other(Cow::Owned(err.to_string())))?;

        Ok(Self {
            include_globset,
            def: cfg_t,
        })
    }
}

pub fn generate_template_output(
    config: &VutConfig,
    root_path: &Path,
    version: &Version,
    dir_entries: &[walkdir::DirEntry],
) -> Result<(), VutError> {
    let specs = build_template_specs(config)?;

    let template_input = generate_template_input(version)?;

    let mut processed_files: Vec<PathBuf> = Vec::new();
    let mut generated_files: Vec<PathBuf> = Vec::new();

    for spec in specs.iter() {
        debug!("{:?}", spec);

        let include_globset = &spec.include_globset;
        let def = spec.def;

        let start_path: Cow<Path> = match &def.start_path {
            Some(p) => Cow::Owned(root_path.join(p)),
            None => Cow::Borrowed(root_path),
        };

        let output_path = match &def.output_path {
            Some(p) => Cow::Owned(root_path.join(p)),
            None => start_path.clone(),
        };

        let processor = def.processor.as_ref().map(|s| s.as_str()).unwrap_or("vut");
        let encoding = def.encoding.as_ref().map(|s| s.as_str());

        let template_files_iter = dir_entries
            .iter()
            .map(|entry| entry.path())
            // Exclude files outside the start path
            .filter(|path| path.starts_with(&start_path))
            // Only include template files
            .filter_map(|path| {
                // Make path relative, as we only want to match on the path
                // relative to the start path.
                let rel_path = path.strip_prefix(&start_path).unwrap();

                if include_globset.is_match(rel_path) {
                    Some((path, rel_path))
                } else {
                    None
                }
            });

        for (path, rel_path) in template_files_iter {
            debug!("Processing template {}", path.display());

            let output_file_name: &OsStr = path.file_stem().unwrap();
            let output_file_path = output_path.join(rel_path.with_file_name(output_file_name));

            template::generate_template_with_processor_name(
                processor,
                path,
                &output_file_path,
                &template_input,
                encoding,
            )
            .map_err(|err| VutError::TemplateGenerate(err))?;

            processed_files.push(path.to_path_buf());
            generated_files.push(output_file_path);
        }
    }

    Ok(())
}

pub fn generate_template_input(version: &Version) -> Result<TemplateInput, VutError> {
    let mut template_input = TemplateInput::new();

    let split_prerelease = version
        .prerelease
        .as_ref()
        .map_or(None, |p| version::split_numbered_prerelease(p));
    let split_build = version
        .build
        .as_ref()
        .map_or(None, |b| version::split_numbered_prerelease(b));

    template_input
        .values
        .insert("FullVersion".to_owned(), version.to_string());
    template_input.values.insert(
        "Version".to_owned(),
        Version {
            build: None,
            ..version.clone()
        }
        .to_string(),
    );
    template_input.values.insert(
        "MajorMinorPatch".to_owned(),
        format!("{}.{}.{}", version.major, version.minor, version.patch),
    );
    template_input
        .values
        .insert("MajorMinor".to_owned(), format!("{}.{}", version.major, version.minor));
    template_input
        .values
        .insert("Major".to_owned(), format!("{}", version.major));
    template_input
        .values
        .insert("Minor".to_owned(), format!("{}", version.minor));
    template_input
        .values
        .insert("Patch".to_owned(), format!("{}", version.patch));
    template_input.values.insert(
        "Prerelease".to_owned(),
        version.prerelease.as_ref().map_or("", |p| p).to_owned(),
    );
    template_input.values.insert(
        "PrereleasePrefix".to_owned(),
        split_prerelease
            .and_then(|sp| Some(sp.0.to_owned()))
            .unwrap_or_else(|| "".to_owned()),
    );
    template_input.values.insert(
        "PrereleaseNumber".to_owned(),
        split_prerelease
            .and_then(|sp| Some(format!("{}", sp.1)))
            .unwrap_or_else(|| "".to_owned()),
    );
    template_input
        .values
        .insert("Build".to_owned(), version.build.as_ref().map_or("", |b| b).to_owned());
    template_input.values.insert(
        "BuildPrefix".to_owned(),
        split_build
            .and_then(|sp| Some(sp.0.to_owned()))
            .unwrap_or_else(|| "".to_owned()),
    );
    template_input.values.insert(
        "BuildNumber".to_owned(),
        split_build
            .and_then(|sp| Some(format!("{}", sp.1)))
            .unwrap_or_else(|| "".to_owned()),
    );

    Ok(template_input)
}

fn build_template_specs(config: &VutConfig) -> Result<Vec<TemplateSpec>, VutError> {
    let mut specs: Vec<TemplateSpec> = Vec::new();

    for cfg_t in config.template.iter() {
        let spec = TemplateSpec::from_config_template(cfg_t)?;

        specs.push(spec);
    }

    Ok(specs)
}
<|endoftext|>


<commit_before>//! Contains implementation of Binary FBX emitter.

extern crate byteorder;
extern crate flate2;

use std::io::{Write, Seek, SeekFrom};
use self::byteorder::{LittleEndian, WriteBytesExt};
use writer::error::{Result, Error};
use common::Property;

/// A writer for Binary FBX.
#[derive(Debug, Clone)]
pub struct BinaryEmitter {
    version: u32,
    pos: u64,
    end_offset_pos_stack: Vec<u64>,
    null_record_necessities: Vec<bool>,
}

impl BinaryEmitter {
    /// Constructs Binary FBX writer with FBX version.
    pub fn new(version: u32) -> Self {
        BinaryEmitter {
            version: version,
            pos: 0,
            end_offset_pos_stack: vec![],
            null_record_necessities: vec![],
        }
    }

    pub fn emit_start_fbx<W: Write + Seek>(&mut self, sink: &mut W, ver: u32) -> Result<()> {
        if (ver < 7000) || (ver >= 8000) {
            error!("Unsupported version: {}", ver);
            return Err(Error::UnsupportedFbxVersion(ver));
        }
        // Write magic binary for Binary FBX.
        try!(sink.write(b"Kaydara FBX Binary  \x00"));
        // Meaning is unknown, but value seems to be always `[0x1A, 0x00]`.
        try!(sink.write(b"\x1a\x00"));
        // Write FBX version.
        try!(sink.write_u32::<LittleEndian>(ver));

        Ok(())
    }

    pub fn emit_end_fbx<W: Write + Seek>(&mut self, sink: &mut W) -> Result<()> {
        // Write null record header.
        // 13: size of a node record header.
        try!(sink.write_all(&[0; 13]));

        // Write footer.

        // Write unknown footer.
        // NOTE: This footer is `fa bc ax 0x dx cx dx 6x bx 7x fx 8x 1x fx 2x 7x`,
        //       but detail is unknown.
        try!(sink.write_all(&[
           0xfa as u8, 0xbc, 0xaf, 0x0f,
           0xdf, 0xcf, 0xdf, 0x6f,
           0xbf, 0x7f, 0xff, 0x8f,
           0x1f, 0xff, 0x2f, 0x7f
        ]));
        // Write padding.
        {
            let current_off = try!(sink.seek(SeekFrom::Current(0))) & 0x0f;
            if current_off != 0 {
                try!(sink.write_all(&(current_off..16).map(|_| 0).collect::<Vec<u8>>()));
            }
        }
        // Write `0u32`, FBX version, and [0; 120].
        try!(sink.write_all(&[0; 4]));
        try!(sink.write_u32::<LittleEndian>(self.version));
        try!(sink.write_all(&[0; 120]));
        // Write unknown but fixed magic.
        try!(sink.write_all(&[
            0xf8 as u8, 0x5a, 0x8c, 0x6a,
            0xde, 0xf5, 0xd9, 0x7e,
            0xec, 0xe9, 0x0c, 0xe3,
            0x75, 0x8f, 0x29, 0x0b
        ]));

        // All done.
        Ok(())
    }

    pub fn emit_start_node<W: Write + Seek>(&mut self, sink: &mut W, name: &str, properties: &[Property]) -> Result<()> {
        if let Some(top) = self.null_record_necessities.last_mut() {
            // Parent node requires null record, because it has child node (the current node!).
            *top = true;
        }
        self.null_record_necessities.push(properties.is_empty());

        // Write node record header.
        // For detail of node record header, see `reader::parser::binary::NodeRecordHeader` struct.
        // Write a placeholder for `end_offset` and remember current offset.
        self.end_offset_pos_stack.push(try!(sink.seek(SeekFrom::Current(0))));
        try!(sink.write_u32::<LittleEndian>(0xef_be_ad_de));
        // Write `num_properties`.
        try!(sink.write_u32::<LittleEndian>(properties.len() as u32));
        // Write a default value of `property_list_len`.
        let prop_list_len_offset = try!(sink.seek(SeekFrom::Current(0)));
        try!(sink.write_u32::<LittleEndian>(0));
        // Write length of the node name.
        try!(sink.write_u8(name.len() as u8));

        // Write a node name.
        try!(sink.write_all(name.as_bytes()));

        // Write properties.
        if !properties.is_empty() {
            let mut props_byte_size = 0_u32;
            for prop in properties {
                macro_rules! read_array_value {
                    ($vec:ident, $type_code:expr, $elem_type_writer:ident) => ({
                        try!(sink.write_u8($type_code as u8));

                        // Write a property array header.
                        // Write array length (element numbers, not byte size).
                        try!(sink.write_u32::<LittleEndian>($vec.len() as u32));
                        // Write encoding.
                        // 0 for plain data, 1 for zlib-compressed data.
                        try!(sink.write_u32::<LittleEndian>(1));
                        // Write a placeholder for byte size of properties.
                        let byte_size_pos = try!(sink.seek(SeekFrom::Current(0)));
                        try!(sink.write_u32::<LittleEndian>(0));

                        let vec_start_pos = try!(sink.seek(SeekFrom::Current(0)));
                        {
                            let mut encoder = flate2::write::ZlibEncoder::new(sink.by_ref(), flate2::Compression::Default);
                            for &v in $vec {
                                //try!(encoder.write_i32::<LittleEndian>(v));
                                try!(encoder.$elem_type_writer::<LittleEndian>(v));
                            }
                            try!(encoder.finish());
                        }
                        let last_pos = try!(sink.seek(SeekFrom::Current(0)));

                        // Update byte size of properties.
                        let byte_size = last_pos - vec_start_pos;
                        try!(sink.seek(SeekFrom::Start(byte_size_pos)));
                        try!(sink.write_u32::<LittleEndian>(byte_size as u32));
                        try!(sink.seek(SeekFrom::Start(last_pos)));
                        // 12: property array header.
                        12 + byte_size as u32
                    })
                };
                props_byte_size += 1 + match *prop {
                    Property::Bool(v) => {
                        try!(sink.write_u8('C' as u8));
                        // `'Y'` is `0x59`,  `'T'` is `0x54`.
                        try!(sink.write_u8(if v { 'Y' } else { 'T' } as u8));
                        1
                    },
                    Property::I16(v) => {
                        try!(sink.write_u8('Y' as u8));
                        try!(sink.write_i16::<LittleEndian>(v));
                        2
                    },
                    Property::I32(v) => {
                        try!(sink.write_u8('I' as u8));
                        try!(sink.write_i32::<LittleEndian>(v));
                        4
                    },
                    Property::I64(v) => {
                        try!(sink.write_u8('L' as u8));
                        try!(sink.write_i64::<LittleEndian>(v));
                        8
                    },
                    Property::F32(v) => {
                        try!(sink.write_u8('F' as u8));
                        try!(sink.write_f32::<LittleEndian>(v));
                        4
                    },
                    Property::F64(v) => {
                        try!(sink.write_u8('D' as u8));
                        try!(sink.write_f64::<LittleEndian>(v));
                        8
                    },
                    Property::VecBool(vec) => {
                        try!(sink.write_u8('b' as u8));
                        for v in vec.iter().map(|&v| if v { 'Y' } else { 'T' } as u8) {
                            try!(sink.write_u8(v));
                        }
                        vec.len() as u32
                    },
                    Property::VecI32(vec) => {
                        read_array_value!(vec, 'i', write_i32)
                    },
                    Property::VecI64(vec) => {
                        read_array_value!(vec, 'l', write_i64)
                    },
                    Property::VecF32(vec) => {
                        read_array_value!(vec, 'f', write_f32)
                    },
                    Property::VecF64(vec) => {
                        read_array_value!(vec, 'd', write_f64)
                    },
                    Property::String(s) => {
                        try!(sink.write_u8('S' as u8));
                        try!(sink.write_u32::<LittleEndian>(s.len() as u32));
                        try!(sink.write_all(s.as_bytes()));
                        4 + s.len() as u32
                    },
                    Property::Binary(b) => {
                        try!(sink.write_u8('R' as u8));
                        try!(sink.write_u32::<LittleEndian>(b.len() as u32));
                        try!(sink.write_all(b));
                        4 + b.len() as u32
                    },
                };
            }
            // Update `property_list_len`
            let last_pos = try!(sink.seek(SeekFrom::Current(0)));
            try!(sink.seek(SeekFrom::Start(prop_list_len_offset)));
            try!(sink.write_u32::<LittleEndian>(props_byte_size));
            try!(sink.seek(SeekFrom::Start(last_pos)));
        }

        Ok(())
    }

    pub fn emit_end_node<W: Write + Seek>(&mut self, sink: &mut W) -> Result<()> {
        // Write a null record header if necessary.
        if let Some(required) = self.null_record_necessities.pop() {
            if required {
                // 13: size of a node record header.
                try!(sink.write_all(&[0; 13]));
            }
        } else {
            return Err(Error::ExtraEndNode);
        }

        // Update `end_offset`.
        let last_pos = try!(sink.seek(SeekFrom::Current(0)));
        try!(sink.seek(SeekFrom::Start(self.end_offset_pos_stack.pop().unwrap())));
        try!(sink.write_u32::<LittleEndian>(last_pos as u32));
        try!(sink.seek(SeekFrom::Start(last_pos)));

        Ok(())
    }
}
<commit_msg>Support FBX 7.5.0 binary<commit_after>//! Contains implementation of Binary FBX emitter.

extern crate byteorder;
extern crate flate2;

use std::io::{Write, Seek, SeekFrom};
use self::byteorder::{LittleEndian, WriteBytesExt};
use writer::error::{Result, Error};
use common::Property;

/// A writer for Binary FBX.
#[derive(Debug, Clone)]
pub struct BinaryEmitter {
    version: u32,
    pos: u64,
    end_offset_pos_stack: Vec<u64>,
    null_record_necessities: Vec<bool>,
}

impl BinaryEmitter {
    /// Constructs Binary FBX writer with FBX version.
    pub fn new(version: u32) -> Self {
        BinaryEmitter {
            version: version,
            pos: 0,
            end_offset_pos_stack: vec![],
            null_record_necessities: vec![],
        }
    }

    pub fn emit_start_fbx<W: Write + Seek>(&mut self, sink: &mut W, ver: u32) -> Result<()> {
        if (ver < 7000) || (ver >= 8000) {
            error!("Unsupported version: {}", ver);
            return Err(Error::UnsupportedFbxVersion(ver));
        }
        // Write magic binary for Binary FBX.
        try!(sink.write(b"Kaydara FBX Binary  \x00"));
        // Meaning is unknown, but value seems to be always `[0x1A, 0x00]`.
        try!(sink.write(b"\x1a\x00"));
        // Write FBX version.
        try!(sink.write_u32::<LittleEndian>(ver));

        Ok(())
    }

    pub fn emit_end_fbx<W: Write + Seek>(&mut self, sink: &mut W) -> Result<()> {
        // Write null record header.
        if self.version < 7500 {
            // 13: size of a node record header (4+4+4+1).
            try!(sink.write_all(&[0; 13]));
        } else {
            // 25: size of a node record header (8+8+8+1).
            try!(sink.write_all(&[0; 25]));
        }

        // Write footer.

        // Write unknown footer.
        // NOTE: This footer is `fa bc ax 0x dx cx dx 6x bx 7x fx 8x 1x fx 2x 7x`,
        //       but detail is unknown.
        try!(sink.write_all(&[
           0xfa as u8, 0xbc, 0xaf, 0x0f,
           0xdf, 0xcf, 0xdf, 0x6f,
           0xbf, 0x7f, 0xff, 0x8f,
           0x1f, 0xff, 0x2f, 0x7f
        ]));
        // Write padding.
        {
            let current_off = try!(sink.seek(SeekFrom::Current(0))) & 0x0f;
            if current_off != 0 {
                try!(sink.write_all(&(current_off..16).map(|_| 0).collect::<Vec<u8>>()));
            }
        }
        // Write `0u32`, FBX version, and [0; 120].
        try!(sink.write_all(&[0; 4]));
        try!(sink.write_u32::<LittleEndian>(self.version));
        try!(sink.write_all(&[0; 120]));
        // Write unknown but fixed magic.
        try!(sink.write_all(&[
            0xf8 as u8, 0x5a, 0x8c, 0x6a,
            0xde, 0xf5, 0xd9, 0x7e,
            0xec, 0xe9, 0x0c, 0xe3,
            0x75, 0x8f, 0x29, 0x0b
        ]));

        // All done.
        Ok(())
    }

    pub fn emit_start_node<W: Write + Seek>(&mut self, sink: &mut W, name: &str, properties: &[Property]) -> Result<()> {
        if let Some(top) = self.null_record_necessities.last_mut() {
            // Parent node requires null record, because it has child node (the current node!).
            *top = true;
        }
        self.null_record_necessities.push(properties.is_empty());

        // Write node record header.
        // For detail of node record header, see `reader::parser::binary::NodeRecordHeader` struct.
        let prop_list_len_offset;
        if self.version < 7500 {
            // Write a placeholder for `end_offset` and remember current offset.
            self.end_offset_pos_stack.push(try!(sink.seek(SeekFrom::Current(0))));
            try!(sink.write_u32::<LittleEndian>(0xef_be_ad_de));
            // Write `num_properties`.
            if properties.len() > u32::max_value() as usize {
                return Err(Error::DataTooLarge(format!("Number of node properties ({}) is too large for FBX {}", properties.len(), self.version)));
            }
            try!(sink.write_u32::<LittleEndian>(properties.len() as u32));
            // Write a default value of `property_list_len`.
            prop_list_len_offset = try!(sink.seek(SeekFrom::Current(0)));
            try!(sink.write_u32::<LittleEndian>(0));
        } else {
            // Write a placeholder for `end_offset` and remember current offset.
            self.end_offset_pos_stack.push(try!(sink.seek(SeekFrom::Current(0))));
            try!(sink.write_u64::<LittleEndian>(0xef_be_ad_de_ef_be_ad_de));
            // Write `num_properties`.
            if properties.len() > u64::max_value() as usize {
                return Err(Error::DataTooLarge(format!("Number of node properties ({}) is too large for FBX {}", properties.len(), self.version)));
            }
            try!(sink.write_u64::<LittleEndian>(properties.len() as u64));
            // Write a default value of `property_list_len`.
            prop_list_len_offset = try!(sink.seek(SeekFrom::Current(0)));
            try!(sink.write_u64::<LittleEndian>(0));
        }
        // Write length of the node name.
        try!(sink.write_u8(name.len() as u8));

        // Write a node name.
        try!(sink.write_all(name.as_bytes()));

        // Write properties.
        if !properties.is_empty() {
            let mut props_byte_size = 0_u64;
            for prop in properties {
                macro_rules! read_array_value {
                    ($vec:ident, $type_code:expr, $elem_type_writer:ident) => ({
                        try!(sink.write_u8($type_code as u8));

                        // Write a property array header.
                        // Write array length (element numbers, not byte size).
                        try!(sink.write_u32::<LittleEndian>($vec.len() as u32));
                        // Write encoding.
                        // 0 for plain data, 1 for zlib-compressed data.
                        try!(sink.write_u32::<LittleEndian>(1));
                        // Write a placeholder for byte size of properties.
                        let byte_size_pos = try!(sink.seek(SeekFrom::Current(0)));
                        try!(sink.write_u32::<LittleEndian>(0));

                        let vec_start_pos = try!(sink.seek(SeekFrom::Current(0)));
                        {
                            let mut encoder = flate2::write::ZlibEncoder::new(sink.by_ref(), flate2::Compression::Default);
                            for &v in $vec {
                                //try!(encoder.write_i32::<LittleEndian>(v));
                                try!(encoder.$elem_type_writer::<LittleEndian>(v));
                            }
                            try!(encoder.finish());
                        }
                        let last_pos = try!(sink.seek(SeekFrom::Current(0)));

                        // Update byte size of properties.
                        let byte_size = last_pos - vec_start_pos;
                        try!(sink.seek(SeekFrom::Start(byte_size_pos)));
                        try!(sink.write_u32::<LittleEndian>(byte_size as u32));
                        try!(sink.seek(SeekFrom::Start(last_pos)));
                        // 12: property array header.
                        12 + byte_size as u64
                    })
                };
                props_byte_size += 1 + match *prop {
                    Property::Bool(v) => {
                        try!(sink.write_u8('C' as u8));
                        // `'Y'` is `0x59`,  `'T'` is `0x54`.
                        try!(sink.write_u8(if v { 'Y' } else { 'T' } as u8));
                        1
                    },
                    Property::I16(v) => {
                        try!(sink.write_u8('Y' as u8));
                        try!(sink.write_i16::<LittleEndian>(v));
                        2
                    },
                    Property::I32(v) => {
                        try!(sink.write_u8('I' as u8));
                        try!(sink.write_i32::<LittleEndian>(v));
                        4
                    },
                    Property::I64(v) => {
                        try!(sink.write_u8('L' as u8));
                        try!(sink.write_i64::<LittleEndian>(v));
                        8
                    },
                    Property::F32(v) => {
                        try!(sink.write_u8('F' as u8));
                        try!(sink.write_f32::<LittleEndian>(v));
                        4
                    },
                    Property::F64(v) => {
                        try!(sink.write_u8('D' as u8));
                        try!(sink.write_f64::<LittleEndian>(v));
                        8
                    },
                    Property::VecBool(vec) => {
                        try!(sink.write_u8('b' as u8));
                        for v in vec.iter().map(|&v| if v { 'Y' } else { 'T' } as u8) {
                            try!(sink.write_u8(v));
                        }
                        vec.len() as u64
                    },
                    Property::VecI32(vec) => {
                        read_array_value!(vec, 'i', write_i32)
                    },
                    Property::VecI64(vec) => {
                        read_array_value!(vec, 'l', write_i64)
                    },
                    Property::VecF32(vec) => {
                        read_array_value!(vec, 'f', write_f32)
                    },
                    Property::VecF64(vec) => {
                        read_array_value!(vec, 'd', write_f64)
                    },
                    Property::String(s) => {
                        try!(sink.write_u8('S' as u8));
                        try!(sink.write_u32::<LittleEndian>(s.len() as u32));
                        try!(sink.write_all(s.as_bytes()));
                        4 + s.len() as u64
                    },
                    Property::Binary(b) => {
                        try!(sink.write_u8('R' as u8));
                        try!(sink.write_u32::<LittleEndian>(b.len() as u32));
                        try!(sink.write_all(b));
                        4 + b.len() as u64
                    },
                };
            }
            // Update `property_list_len`
            let last_pos = try!(sink.seek(SeekFrom::Current(0)));
            try!(sink.seek(SeekFrom::Start(prop_list_len_offset)));
            if self.version < 7500 {
                if props_byte_size > u32::max_value() as u64 {
                    return Err(Error::DataTooLarge(format!("Properties size ({} bytes) is too large for FBX {}", props_byte_size, self.version)));
                }
                try!(sink.write_u32::<LittleEndian>(props_byte_size as u32));
            } else {
                try!(sink.write_u64::<LittleEndian>(props_byte_size));
            }
            try!(sink.seek(SeekFrom::Start(last_pos)));
        }

        Ok(())
    }

    pub fn emit_end_node<W: Write + Seek>(&mut self, sink: &mut W) -> Result<()> {
        // Write a null record header if necessary.
        if let Some(required) = self.null_record_necessities.pop() {
            if required {
                if self.version < 7500 {
                    // 13: size of a node record header (4+4+4+1).
                    try!(sink.write_all(&[0; 13]));
                } else {
                    // 25: size of a node record header (8+8+8+1).
                    try!(sink.write_all(&[0; 25]));
                }
            }
        } else {
            return Err(Error::ExtraEndNode);
        }

        // Update `end_offset`.
        let last_pos = try!(sink.seek(SeekFrom::Current(0)));
        try!(sink.seek(SeekFrom::Start(self.end_offset_pos_stack.pop().unwrap())));
        if self.version < 7500 {
            if last_pos > u32::max_value() as u64 {
                return Err(Error::DataTooLarge(format!("File size (currently {} bytes) is too large for FBX {}", last_pos, self.version)));
            }
            try!(sink.write_u32::<LittleEndian>(last_pos as u32));
        } else {
            try!(sink.write_u64::<LittleEndian>(last_pos));
        }
        try!(sink.seek(SeekFrom::Start(last_pos)));

        Ok(())
    }
}
<|endoftext|>