Zenn
🔗

URL短縮サービスを勢いだけで実装してみる（Rust, axum）

2022/09/03に公開
Rust
#
axum
tech

https://www.shuttle.rs/blog/2022/03/13/url-shortener
こちらのノリノリの記事にインスパイアされたので、axumを使ったURL短縮サービスを実装してみたいと思います（※バックエンドだけ）。

めちゃくちゃ雑に考えると、URL短縮サービスとは、

  • URLを受け取り、短縮されたURLを返す
  • その短縮版URLにアクセスすると、元のURLにリダイレクトされる

というものです。
したがって、短縮URLをkey、元のURLをvalueとしてストアするデータベースが必要になります。
丁寧にやっていくならpostgresなどの永続するデータベースで実装しなければなりませんが、今回はsledを使ってみることにします。
sledはRustで書かれたインメモリのkey/value型のデータベースです。

sled is a high-performance embedded database with an API that is similar to a BTreeMap<[u8], [u8]>, but with several additional capabilities for assisting creators of stateful systems.
https://docs.rs/sled/latest/sled/index.html

さらにすべての操作がatomicということで、インメモリということを気にしなければサーバーサイドでも活用できそうです。

hello, world

さて、まずはaxumのhello, worldをとにかく形にします。

[dependencies]
axum = {version = "0.5.15", features = ["macros"]}
env_logger = "0.9.0"
log = "0.4.17"
tokio = {version = "1.21.0", features = ["full"]}

//main.rs

mod router;

use axum::{routing::get, Router};
use router::*;

#[tokio::main]
async fn main() {
    env_logger::init();
    println!("Server started.");

    let app = Router::new().route("/", get(health));

    axum::Server::bind(&"0.0.0.0:8080".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

//router.rs

use axum::debug_handler;

#[debug_handler]
pub async fn health() -> String {
    "Hello, developer.".to_string()
}

# cargo runの後...
$ xh -b :8080
Hello, developer.

無事ヘルスチェックができました。ちなみに#[debug_handler]はハンドラ関数のエラーメッセージをわかりやすくしてくれるマクロです。
注意点として、cargo add tokioしただけでは#[tokio::main]がうまく動いてくれませんので、features=["macros"]としておきましょう。

Rustで一番楽しいのはエラーハンドリングだと思っている私ですが、今回はシンプルに進めていくためthiserrorもanyhowもナシでいきます。

UUIDの生成

次に、URLから短縮形を返すハンドラを実装していきます。
各URLに対し固有の短縮形を返さないといけないので、とにかくまずuuidを生成して返す関数を作りましょう。今回は元記事と同じく、nanoidを使ってシンプルに作ります。

//main.rs

    let app = Router::new()
        .route("/", get(health))
        .route("/shorten", post(shorten));

//router.rs

#[debug_handler]
pub async fn shorten(body: String) -> String {
    let uuid = nanoid::nanoid!(8);
    return uuid;
}

$ xh post localhost:8080/shorten
-0rZCBia

リクエストに関係なく、8桁のUUIDを返すものを実装しました。

データベースへ登録

次に、このUUIDと、リクエストの中身のURLを紐付けて、データベースへ保存します。
データベースの初期化はmain関数の中で行い、それをshared stateとしてハンドラと共有するためにaxum::Extensionを使います。

//main.rs

#[tokio::main]
async fn main() {
    env_logger::init();
    info!("Server started.");

    let db: sled::Db = sled::open("my_db").unwrap();
    info!("Database started.");

    let app = Router::new()
        .route("/", get(health))
        .route("/shorten", post(shorten))
        .layer(Extension(db)); // shared stateとしてハンドラと共有する

    axum::Server::bind(&"0.0.0.0:8080".parse().unwrap())
        .serve(app.into_make_service())
        .await
        .unwrap();
}

ハンドラ側でshared stateを呼ぶには以下のようにします。

//router.rs

// データベースに登録するため、JSONを受け取る形にします。
#[derive(Debug, Deserialize)]
pub struct Payload { 
    url: String,
} 

#[debug_handler]
pub async fn shorten(Json(payload): Json<Payload>, Extension(db): Extension<sled::Db>) -> String {
    info!("{:?}", payload);
    let mut uuid = nanoid::nanoid!(8);
    // uuidが登録済のものと被らないか確認
    while db.contains_key(&uuid).unwrap() { 
        uuid = nanoid::nanoid!(8);
    } 
    let url_as_bytes = payload.url.as_bytes();
    db.insert(&uuid, url_as_bytes).unwrap();
    info!("key: {}, value: {:?}", uuid, url_as_bytes);
    assert_eq!(&db.get(uuid.as_bytes()).unwrap().unwrap(), url_as_bytes);
    uuid
}

# sample.json
{"url": "https://google.com"}

$ xh -b :8080/shorten < sample.json
ZAbBi7Hw

パニックを起こさないので問題なくDBへ追加できているようです。

データベースから引き出す

続けて、DBからvalueを引き出すための関数…/redirect/:idへアクセスして元のURLを返すハンドラを実装していきましょう。

//main.rs
//main内でURLをキャプチャする
    let app = Router::new()
        .route("/", get(health))
        .route("/shorten", post(shorten))
        .route("/redirect/:id", get(redirect))
        .layer(Extension(db));

shorten関数と同じく、redirect関数内でもExtension(db)でshared stateを取り出していきます。

//router.rs

#[debug_handler]
pub async fn redirect(Path(id): Path<String>, Extension(db): Extension<sled::Db>) -> String {
    match &db.get(&id).unwrap() {
        Some(url) => {
            let url = String::from_utf8(url.to_vec()).unwrap();
            info!("URL found: {:#?}", url);
            url
        }
        None => "Error: Not found.".to_string(),
    }
}

$ xh post :8080/shorten < sample.json
-8ij5irt

$ xh :8080/redirect/-8ij5irt
https://google.com

これでデータベースへの出し入れが実装できました。

リダイレクト

あとはリダイレクトの実装だけです。redirect関数の返り値を、axum::Redirectとします。データベースにkeyがなかった場合はとりあえず/へリダイレクトしています。

//router.rs

#[debug_handler]
pub async fn redirect(Path(id): Path<String>, Extension(db): Extension<sled::Db>) -> Redirect {
    match &db.get(&id).unwrap() {
        Some(url) => {
            let url = String::from_utf8(url.to_vec()).unwrap();
            info!("URL found: {:#?}", url);
            Redirect::to(&url)
        }
        None => {
            info!("URL not found.");
            Redirect::to("/")
        }
    }
}

$ xh -b :8080/shorten < sample.json
9B3LgnGc

xh :8080/redirect/9B3LgnGc
HTTP/1.1 303 See Other
Content-Length: 0
Date: Sat, 03 Sep 2022 05:23:13 GMT
Location: https://google.com

実際にlocalhost:8080/redirect/9B3LgnGcへブラウザでアクセスしてみると、google.comへリダイレクトされ、ページが表示されました。

実際のサービスで/redirect/hogehogeというURLになるのはダサいので、ルーティングをもうちょっとよしなにできるはずですが、満足したのでいったんここでお開きとさせていただきます。ありがとうございました。

今回のソースコードはこちら。
https://github.com/kyoheiu/url-shortener-axum-sled

Discussion