amqp_worker 0.5.9

AMQP Worker to listen and provide trait to process message
Documentation
#![doc(html_favicon_url = "https://media-io.com/images/mediaio_logo.png")]
#![doc(html_logo_url = "https://media-io.com/images/mediaio_logo.png")]
#![doc(html_no_source)]

extern crate amq_protocol_types;
extern crate amq_protocol_uri;
extern crate chrono;
extern crate env_logger;
extern crate failure;
extern crate futures;
#[macro_use]
extern crate log;
extern crate lapin_futures as lapin;
extern crate reqwest;
extern crate serde;
#[macro_use]
extern crate serde_derive;
#[macro_use]
extern crate serde_json;
extern crate tokio;

mod config;
pub mod job;
mod message;

use amq_protocol_types::AMQPValue;
use amq_protocol_uri::*;
use chrono::prelude::*;
use config::*;
use env_logger::Builder;
use failure::Error;
use futures::{future::Future, Stream};
use job::JobResult;
use lapin::{options::*, types::FieldTable, ConnectionProperties};
use std::{env, fs, io::Write, thread, time};
use tokio::runtime::Runtime;

pub trait MessageEvent {
  fn process(&self, _message: &str) -> Result<JobResult, MessageError>
  where
    Self: std::marker::Sized,
  {
    Err(MessageError::NotImplemented())
  }
}

#[derive(Debug, PartialEq)]
pub enum MessageError {
  RuntimeError(String),
  ProcessingError(JobResult),
  RequirementsError(String),
  NotImplemented(),
}

fn load_docker_container_id(filename: &str) -> String {
  match fs::read_to_string(filename) {
    Ok(content) => parse_docker_container_id(&content),
    Err(_msg) => "unknown".to_string(),
  }
}

fn parse_docker_container_id(content: &str) -> String {
  let lines: Vec<&str> = content.split('\n').collect();
  if lines.is_empty() {
    return "unknown".to_string();
  }
  let items: Vec<&str> = lines[0].split(':').collect();
  if items.len() != 3 {
    return "unknown".to_string();
  }

  let long_identifier: Vec<&str> = items[2].split("/docker/").collect();
  if long_identifier.len() != 2 {
    return "unknown".to_string();
  }
  let mut identifier = long_identifier[1].to_string();
  identifier.truncate(12);
  identifier
}

#[test]
fn test_load_docker_container_id() {
  assert_eq!(
    load_docker_container_id("./tests/cgroup.sample"),
    "da9002cb1553".to_string()
  );

  assert_eq!(
    load_docker_container_id("/tmp/file_not_exists"),
    "unknown".to_string()
  );

  assert_eq!(parse_docker_container_id(""), "unknown".to_string());
  assert_eq!(parse_docker_container_id("\n"), "unknown".to_string());
  assert_eq!(parse_docker_container_id("a:b:c\n"), "unknown".to_string());
}

pub fn start_worker<ME: MessageEvent>(message_event: &'static ME)
where
  ME: std::marker::Sync,
{
  let mut builder = Builder::from_default_env();
  let container_id = load_docker_container_id("/proc/self/cgroup");
  let queue = get_amqp_queue();

  builder
    .format(move |stream, record| {
      writeln!(
        stream,
        "{} - {} - {} - {} - {} - {}",
        Utc::now(),
        &container_id,
        queue,
        record.target().parse::<i64>().unwrap_or(-1),
        record.level(),
        record.args(),
      )
    })
    .init();

  let queue = env::var("AMQP_QUEUE").expect("missing AMQP queue");
  let version = env::var("VERSION").unwrap_or_else(|_| "unknown".to_string());

  info!("Worker: {}, version: {}", queue, version);

  loop {
    let amqp_tls = get_amqp_tls();
    let amqp_hostname = get_amqp_hostname();
    let amqp_port = get_amqp_port();
    let amqp_username = get_amqp_username();
    let amqp_password = get_amqp_password();
    let amqp_vhost = get_amqp_vhost();
    let amqp_queue = get_amqp_queue();
    let amqp_completed_queue = get_amqp_completed_queue();
    let amqp_error_queue = get_amqp_error_queue();

    info!("Start connection with configuration:");
    info!("AMQP TLS: {}", amqp_tls);
    info!("AMQP HOSTNAME: {}", amqp_hostname);
    info!("AMQP PORT: {}", amqp_port);
    info!("AMQP USERNAME: {}", amqp_username);
    info!("AMQP VHOST: {}", amqp_vhost);
    info!("AMQP QUEUE: {}", amqp_queue);

    let scheme = if amqp_tls {
      AMQPScheme::AMQPS
    } else {
      AMQPScheme::AMQP
    };

    let amqp_uri = AMQPUri {
      scheme,
      authority: AMQPAuthority {
        userinfo: AMQPUserInfo {
          username: amqp_username,
          password: amqp_password,
        },
        host: amqp_hostname,
        port: amqp_port,
      },
      vhost: amqp_vhost,
      query: Default::default(),
    };

    let state = Runtime::new().unwrap().block_on_all(
      lapin::Client::connect_uri(amqp_uri, ConnectionProperties::default())
        .map_err(Error::from)
        .and_then(|client| client.create_channel().map_err(Error::from))
        .and_then(move |channel| {
          let id = channel.id();
          debug!("created channel with id: {}", id);

          let prefetch_count = 1;
          if let Err(msg) = channel
            .basic_qos(prefetch_count, BasicQosOptions::default())
            .wait()
          {
            error!("Unable to set QoS on channels: {:?}", msg);
          }

          let ch = channel.clone();

          let delayed_name = amqp_queue.clone() + "_delayed";
          let exchange_type = "fanout";

          if let Err(msg) = channel
            .exchange_declare(
              &delayed_name,
              exchange_type,
              ExchangeDeclareOptions::default(),
              FieldTable::default(),
            )
            .wait()
          {
            error!("Unable to create exchange {}: {:?}", delayed_name, msg);
          }

          let mut delaying_queue_fields = FieldTable::default();
          delaying_queue_fields.insert(
            "x-dead-letter-exchange".into(),
            AMQPValue::LongString("".into()),
          );
          delaying_queue_fields.insert("x-message-ttl".into(), AMQPValue::ShortInt(5000));

          if let Err(msg) = channel
            .queue_declare(
              &delayed_name,
              QueueDeclareOptions::default(),
              delaying_queue_fields,
            )
            .wait()
          {
            error!("Unable to create queue {}: {:?}", delayed_name, msg);
          }

          let routing_key = "*";

          if let Err(msg) = channel
            .queue_bind(
              &delayed_name,
              &delayed_name,
              routing_key,
              QueueBindOptions::default(),
              FieldTable::default(),
            )
            .wait()
          {
            error!("Unable to create queue {}: {:?}", delayed_name, msg);
          }

          if let Err(msg) = channel
            .queue_declare(
              &amqp_completed_queue,
              QueueDeclareOptions::default(),
              FieldTable::default(),
            )
            .wait()
          {
            error!("Unable to create queue {}: {:?}", amqp_completed_queue, msg);
          }

          if let Err(msg) = channel
            .queue_declare(
              &amqp_error_queue,
              QueueDeclareOptions::default(),
              FieldTable::default(),
            )
            .wait()
          {
            error!("Unable to create queue {}: {:?}", amqp_error_queue, msg);
          }

          let mut queue_fields = FieldTable::default();
          queue_fields.insert(
            "x-dead-letter-exchange".into(),
            AMQPValue::LongString(delayed_name.into()),
          );
          queue_fields.insert(
            "x-dead-letter-routing-key".into(),
            AMQPValue::LongString(amqp_queue.clone().into()),
          );

          channel
            .queue_declare(&amqp_queue, QueueDeclareOptions::default(), queue_fields)
            .and_then(move |queue| {
              info!("channel {} declared queue {}", id, amqp_queue);

              channel.basic_consume(
                &queue,
                "amqp_worker",
                BasicConsumeOptions::default(),
                FieldTable::default(),
              )
            })
            .and_then(move |stream| {
              // process_stream(message_event, channel, stream);
              warn!("start listening stream");

              stream.for_each(move |message| {
                trace!("raw message: {:?}", message);
                message::process_message(message_event, message, &ch);
                Ok(())
              })
            })
            .map_err(Error::from)
        })
        .map_err(Error::from),
    );

    warn!("{:?}", state);
    let sleep_duration = time::Duration::new(1, 0);
    thread::sleep(sleep_duration);
  }
}

#[test]
fn empty_message_event_impl() {
  #[derive(Debug)]
  struct CustomEvent {}

  impl MessageEvent for CustomEvent {}

  let custom_event = CustomEvent {};

  let result = custom_event.process("test");
  assert!(result == Err(MessageError::NotImplemented()));
}