RustでRailsの代わりにCSVを作る

Rustを勉強してて何か実際に使ってみたかった頃に、RailsなWebアプリケーションでCSVをダウンロードする仕組みがあって、普通に作ってみたところメモリを食いつぶして動かなかったので、Rustで作ってみました。

Railsでは以下のように行っていました。

1. 作成するCSVに関するパラメーター取得
2. DBからSELECT
3. CSVファイル作成
4. ZIP圧縮
5. レスポンスで返す

ここで2,3をRustで実行するようにしました。
さらにRailsを楽にするため
4は外部コマンドで実行
5はX-Accel-Redirectを使ってnginxに任せました。

CSVの出力範囲は1日を想定していて
Railsからはffi経由で呼び出しています。

結果、14万件のデーターが6分くらいでダウンロードできました。ZIPファイルのサイズは27MBでした。
メモリは10MBくらい使っていました。
Railsだと4GBメモリを積んだマシンでメモリ食いつぶしてサーバーごとお亡くなりだったので、Rustの威力を見せつけられました。
RailsとRustを組み合わせると効率的なアプリケーションが組めると実感しました。

Rust1.13.0で作りました。

Cargo.toml

[package]
name = "make_csv"
version = "0.1.0"
authors = ["aoyagikouhei"]

[lib]
name = "make_csv"
crate-type = ["dylib"]

[dependencies]
libc = "*"
chrono = "*"
postgres = { version="*", features=["with-chrono"]}

lib.rs

extern crate libc;
extern crate postgres;
extern crate chrono;

mod csv;

use std::ffi::CStr;

#[no_mangle]
pub extern fn execute(
    connection: *const libc::c_char, 
    file_path: *const libc::c_char, 
    start_at: *const libc::c_char, 
    end_at: *const libc::c_char
){
    let connection_cstr = unsafe { CStr::from_ptr(connection) };
    let file_path_cstr = unsafe { CStr::from_ptr(file_path) };
    let start_at_cstr = unsafe { CStr::from_ptr(start_at) };
    let end_at_cstr = unsafe { CStr::from_ptr(end_at) };
    let connection_str = connection_cstr.to_str().unwrap();
    let file_path_str = file_path_cstr.to_str().unwrap();
    let start_at_str = start_at_cstr.to_str().unwrap();
    let end_at_str = end_at_cstr.to_str().unwrap();

    let mut maker = csv::CsvMaker::new(
        connection_str,
        file_path_str,
        start_at_str,
        end_at_str
    );
    maker.execute();
}

csv.rs

extern crate postgres;
extern crate chrono;

use postgres::{Connection, TlsMode};
use std::io::prelude::*;
use std::fs::File;
use chrono::{DateTime, Local, Duration};
use std::cell::RefCell;

fn make_sql() -> String {
    let sql = r#"
        SELECT
            COALESCE(t1.tweet_json->'user'->>'screen_name', '')
            ,COALESCE(t1.tweet_json->>'text', '')
            ,COALESCE(to_char(
                to_timestamp(
                    t1.tweet_json->>'created_at'
                    ,'Dy Mon DD HH24:MI:SS +0000 YYYY'
                )::TIMESTAMPTZ + interval'9 hour'
                ,'YYYY/MM/DD HH24:MI'
            ), '')
            ,COALESCE(t1.tweet_json->'user'->>'statuses_count', '')
            ,COALESCE(t1.tweet_json->'user'->>'friends_count', '')
            ,COALESCE(t1.tweet_json->'user'->>'followers_count', '')
            ,COALESCE(t1.tweet_json->'user'->>'listed_count', '')
            ,COALESCE(t1.tweet_json->'user'->>'description', '')
            ,COALESCE(to_char(
                to_timestamp(
                    t1.tweet_json->'user'->>'created_at'
                    ,'Dy Mon DD HH24:MI:SS +0000 YYYY'
                )::TIMESTAMPTZ + interval'9 hour'
                ,'YYYY/MM/DD HH24:MI'
            ), '')
        FROM
            public.tweets AS t1
        WHERE
            t1.created_at >= $1
            AND t1.created_at < $2
        ORDER BY
            t1.created_at ASC
    "#;
    sql.to_string()
}

pub struct CsvMaker {
    pub connection: Connection,
    pub file: RefCell<File>,
    pub start_at: DateTime<Local>,
    pub end_at: DateTime<Local>,
    pub sql: String
}

impl CsvMaker {
    pub fn new(
        connection_str: &str,
        file_path: &str,
        start_at_str: &str,
        end_at_str: &str
    ) -> CsvMaker {
        let connection = Connection::connect(
            connection_str, TlsMode::None).unwrap();
        let file = RefCell::new(File::create(file_path).unwrap());
        let start_at = start_at_str.parse::<DateTime<Local>>().unwrap();
        let end_at = end_at_str.parse::<DateTime<Local>>().unwrap();
        let sql = make_sql();
        CsvMaker {
            connection: connection,
            file: file,
            start_at: start_at,
            end_at: end_at,
            sql: sql
        }
    }

    pub fn execute(&mut self) {
        let header = format!(
            "\"{}\",\"{}\",\"{}\",\"{}\",\"{}\",\"{}\",\"{}\",\"{}\",\"{}\",\"{}\"\n", 
            "アカウント名", 
            "ツイート内容",
            "投稿日時",
            "ツイート数",
            "フォロー数",
            "フォロワー数",
            "いいね数",
            "リスト数",
            "自己紹介",
            "アカウント登録日時"
            );
        {
            let mut f = self.file.borrow_mut();
            let _ = f.write(b"\xEF\xBB\xBF");
            let _ = f.write(header.as_bytes());
        }

        let mut next_at: DateTime<Local> = self.start_at.clone();

        loop {
            self.execute_one(&next_at);
            next_at = next_at + chrono::Duration::seconds(3600);
            if next_at >= self.end_at {
                break;
            }
        }
        {
            let mut f = self.file.borrow_mut();
            let _ = f.flush();
        }
    }

    fn execute_one(&self, next_at: &DateTime<Local>) {
        let next_end_at: DateTime<Local> = *next_at + Duration::seconds(3600);
        let mut f = self.file.borrow_mut();
        for row in &self.connection.query(&self.sql, &[&next_at, &next_end_at]).unwrap() {
            let screen_name: String = row.get(0);
            let tweet_content: String = row.get(1);
            let tweet_created_at: String = row.get(2);
            let statuses_count: String = row.get(3);
            let friends_count: String = row.get(4);
            let followers_count: String = row.get(5);
            let favourites_count: String = row.get(6);
            let listed_count: String = row.get(7);
            let description: String = row.get(8);
            let user_created_at: String = row.get(9);

            let line = format!(
                "\"{}\",\"{}\",\"{}\",\"{}\",\"{}\",\"{}\",\"{}\",\"{}\",\"{}\",\"{}\"\n", 
                screen_name,
                tweet_content,
                tweet_created_at,
                statuses_count,
                friends_count,
                followers_count,
                favourites_count,
                listed_count,
                description,
                user_created_at
            );
            let _ = f.write(line.as_bytes());
        }
    }
}

csv.rb

require 'ffi'
require 'systemu'

module App 
  class Csv 
    extend FFI::Library
    ffi_lib '../make_csv/target/release/libmake_csv.so'
    attach_function :execute, [:string, :string, :string, :string]

     def initialize()
      Dir.mkdir("/tmp/tweets") if !Dir.exist?("/tmp/tweets")
    end

    def exec(start_at, end_at)
      dir = Dir.mktmpdir("tweets/")
      self.class.execute(
        "postgres://localhost:5432/web_development", 
        dir + "/tweets.csv", 
        start_at, 
        end_at)
      Dir.chdir(dir) do
        systemu "/usr/bin/zip -m tweets.zip tweets.csv"
      end
      dir + "/tweets.zip"
    end
  end
end