Redis モジュールで遊んでみた

ついに Redis がローダブル・モジュールに対応する兆しを見せています。Redis の開発者である antirez 氏のブログ記事 [Redis Loadable Modules System](http://antirez.com/news/106) で 4 ヶ月ほど前に紹介され、[antirez/redis](https://github.com/antirez/redis) の `unstable` ブランチで精力的に(??)開発されています。
遊んでみようと思ってから早 4 ヶ月、いい加減重い腰を上げて今回触ってみたので紹介したいと思います。

## tl; dr

- Redis モジュールで独自データ型作れる
- 今回はトライ木作って単語のオートコンプリートやってみた
    - 永続化やレプリケーションも簡単に対応できた
- 夢が広がる

## Redis モジュールとは

Redis の動作を拡張するためのオブジェクトファイルで、動的にロード/アンロードすることができます。

バージョン違いの Redis に対してモジュールが互換性を保てるように、そして、後方互換性のため Redis コアの開発が足を引っ張られないように、モジュールは Redis コアと「モジュール API」を介してデータをやりとりすることになっています。ちなみに antirez 氏は API の互換性を最重要視していて、今日書いたモジュール/コンパイルしたバイナリが 4 年後もそのまま動く/ロードできるようにしたい、と語っています。背景や思想などについて、さらに詳しいことを知りたい方は、前述の antirez 氏のブログ記事や [antirez/redis レポジトリのイントロダクション](https://github.com/antirez/redis/blob/unstable/src/modules/INTRO.md) (ただしあまりメンテされてない) などを参照してもらえればと思います。

モジュール API についてもまだまだ発展途上で、ドキュメントがまともにメンテされてないのでソースコードを読むのが確実です。インターフェースは [src/redismodule.h](https://github.com/antirez/redis/blob/unstable/src/redismodule.h) にある、`RedisModule_` というプレフィックスで始まる一見関数プロトタイプ宣言に見える関数ポインタ達です。その本体は [src/module.c](https://github.com/antirez/redis/blob/unstable/src/module.c) にある、`RM_` というプレフィックスで始まる、同一のサフィックスをもつ関数になります。これらの関数ポインタは Redis モジュールの読み込み時に呼ばれる `RedisModule_Init` (また後で触れます) にてセットされます。こういう設計にしたのは `RedisModule_Init` 関数に渡すモジュール API のバージョンに応じて関数本体を切り替えるためだと思います。たぶん。ざーっと眺めてみると、`RedisModule_ListPush`、`RedisModule_ZsetFirstInScoreRange` とか `RedisModule_ReplyWithArray` みたいなそれっぽい関数(ポインタ)名があったり、`RedisModule_CreateDataType` みたいな心ときめくものも見つけることができます。

## 簡単な Redis モジュール

とりあえず、あるリストにランダムな文字を追加していくだけの `hello.hoge` というコマンドを提供するモジュールを作ってみます。イメージとしては、以下のように動きます。

```redis
127.0.0.1:6379> get list
(nil)
127.0.0.1:6379> hello.hoge list
o
127.0.0.1:6379> hello.hoge list
m
127.0.0.1:6379> lrange list 0 -1
1) "o"
2) "m"
```

何の役にも立たなそうなコマンドですが、キーを使った Redis のデータ構造の読み書きをする簡単な例としてはありかな、と思ってこれにしました。

### hello.hoge コマンドの実装

早速ですが、コマンド本体のソースコードがこちらです。

```c:sample.c(前半)
#include <stdlib.h>
#include "redismodule.h"  // おまじない

// hello.hoge コマンドの本体
int HelloHoge_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
    // アリティのチェック
    if (argc != 2) return RedisModule_WrongArity(ctx);

    // キーを read/write で開く
    RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_READ | REDISMODULE_WRITE);

    // ランダム文字作る
    char string[] = "a";
    string[0] += rand() % 26;

    // lpush で文字をリストに入れる
    if (RedisModule_ListPush(key, REDISMODULE_LIST_TAIL, RedisModule_CreateString(ctx, string, 1)) == REDISMODULE_ERR)
      return REDISMODULE_ERR;

    // 生成したランダム文字をクライアントに返す
    if (RedisModule_ReplyWithSimpleString(ctx, string) == REDISMODULE_ERR)
      return REDISMODULE_ERR;

    // キーを閉じる
    RedisModule_CloseKey(key);

    return REDISMODULE_OK;
}
```

ざっくりした処理の流れはインラインコメントを参照していただくとして、以下、細かい部分の説明などを書いてみます。

関数名は `#{コマンド名}_RedisCommand` のように命名されることが多いです。そして、引数は

- モジュールのコンテキストオブジェクト
  - これは Redis サーバが呼び出し時に勝手に渡してくれる
- 任意の個数 (`argc` 個) の文字列
  - クライアントが発行したコマンドとその引数文字列

という形式に統一されています。`hello.hoge` コマンドは引数を 1 個 (リストのキー) 取るので、コマンド名を含めて `argc` が `2` かどうかをチェックしています。

Redis のデータ構造にアクセスする場合は、最初に `RedisModule_OpenKey` という関数でキーを `open` する必要があります。このようになっているのは、主にトランザクション周りの取り扱いのためみたいです (詳しい所まで読めてないですが)。また、キー名から Redis の内部データポインタへのルックアップが一回ですむという利点もあったりしそうです。そして、`open` したキーは `RedisModule_CloseKey` で閉じる必要があります。ただし、自動メモリ管理機能というものを ON にしておくと勝手に閉じたりしてくれるようです (がまだちゃんと調べてない)。

`RedisModule_ListPush` は `lpush` です。第二引数でリストの先頭と末尾どちらに追加するか指定できます。最後に、`RedisModule_ReplyWith*` 系の関数でクライアントへの返り値をセットします。文字列以外にも、NULL、整数、配列 (多次元も可)、少数などを返すことができます。

以上が、`hello.hoge` コマンド本体の実装です。

### `hello.hoge` コマンドの登録

コマンド本体が実装できたので、後はこれをモジュールがロードされた時に Redis に登録するだけです。これはモジュールロード時のエントリーポイント `RedisModule_OnLoad` で行います。

```c:sample.c(後半)
int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
    // モジュールの初期化
    if (RedisModule_Init(ctx, "hello", 1, REDISMODULE_APIVER_1) == REDISMODULE_ERR)
      return REDISMODULE_ERR;

    // hello.hoge コマンドの登録
    if (RedisModule_CreateCommand(ctx, "hello.hoge", HelloHoge_RedisCommand, "write deny-oom random fast", 1, 1, 1) == REDISMODULE_ERR)
      return REDISMODULE_ERR;

    return REDISMODULE_OK;
}
```

最初に、冒頭でちょっと触れた `RedisModule_Init` を呼び出します。引数はコンテキストオブジェクト、モジュール名、モジュールのバージョン、モジュール API のバージョンを指定します。
そして、`RedisModule_CreateCommand` で `hello.hoge` というコマンド名に対してコマンド本体 `HelloHoge_RedisCommand` の関数ポインタを登録します。第 4 引数以降はそれぞれ、コマンドのフラグ、引数のキーの開始位置、終了位置、キーの間隔になっています。フラグは多岐にわたるので、詳しくは [src/module.c](https://github.com/antirez/redis/blob/unstable/src/module.c) の `RM_CreateCommand` のコメントを参照してください。`hello.hoge` ではキーは最初の引数一つだけなので、キーの開始位置と終了位置は共に `1` で、間隔も `1` にしています。キーを引数として一切取らない場合は開始位置を `0` とすることになっています。

### モジュールの使い方

さて、上記のコード `sample.c` をビルドしてみましょう。[redismodule.h](https://github.com/antirez/redis/blob/unstable/src/redismodule.h)をインクルードパスのある場所において、

```sh
gcc -O2 -shared -fPIC sample.c -o libsample.so
```

でビルドできます。

また、Redis サーバがモジュール対応してないと当然使えないので、[antirez/redis](https://github.com/antirez/redis) の unstable ブランチで Redis をビルドしてサーバを立ち上げます。そして、`redis-cli` から `module load` コマンドでモジュールを読み込むと、コマンドが使えるようになります。

```redis
127.0.0.1:6379> module load "/path/to/libsample.so"
OK
127.0.0.1:6379> hello.hoge list
x
```

また、毎回ロードするのが面倒な場合は、`redis.conf` に `loadmodule "/path/to/libsample.so"` と書くとサーバ起動時に自動的に読み込まれます。

## ここからが本編

さて、「Redis モジュール」という言葉を聞いた時、きっと言葉では言い表せない「ときめき」みたいなものを感じた方が多いかと思います。多様なデータ型をもつ、データストアとしてはやや特殊な Redis に新たな独自データ型 (ツリー、グラフ、転置インデックスなどなど) を追加できるのではないかというときめき！

というわけで、この章では実際に新たなデータ型を定義してみようと思います。

### 独自データ型を扱うためのモジュール API

まず、独自データ型を扱うためのモジュール API について一通り簡単に説明します。

#### `RedisModule_CreateDataType`

新しいデータ型を定義する関数で、モジュールのエントリポイント `RedisModule_OnLoad` で呼び出します。以下を引数に取ります。

- 型の識別子 (何故か 9 文字固定)。`TYPE` コマンドで返ってきたりする値
- モジュールのバージョン
- RDB からデータをロードする関数
- データを RDB にセーブする関数
- データを AOF に書き出す関数
- `DEBUG DIGEST` コマンド用のダイジェスト関数
- データのメモリ領域を開放する関数

関数の返り値として、`RedisModuleType` を返してきます。この返り値はキーの型チェックや、キーに値をセットする際に使うため、グローバル変数などに格納しておきます。

#### `RedisModule_ModuleTypeGetType`

あるキーの `RedisModuleType` を返します。前述のキーの型チェックに使います。

#### `RedisModule_ModuleTypeSetValue`, `RedisModule_ModuleTypeGetValue`

あるキーに対して、定義した型のデータを格納したり、取得するための関数です。

### トライ木

今回実装してみるデータ型は[トライ木](https://ja.wikipedia.org/wiki/%E3%83%88%E3%83%A9%E3%82%A4%E6%9C%A8)で、ひとまず単語のオートコンプリート機能を実現することを目標とします。トライ木については説明しないので、詳しくはリンク先の Wikipedia を参照してください。

ここでは簡単のため、トライ木のキー (格納する単語) は英小文字だけで構成され、キーに紐付いた値は特にないものとします。というわけで、トライ木のノードは以下の様な構造体になります。

```c:hello_trie.c
typedef struct TrieTypeNode {
    uint8_t terminal;
    struct TrieTypeNode* children[26];
} TrieTypeNode;
```

`terminal` がキーの終端を表し、`children` が各英小文字の対応する枝に繋がった子ノードです。

この構造体に対して、`RedisModule_CreateDataType` に渡す各種関数を実装していきます。まずは、RDB からデータをロードする関数です。

```c:hello_trie.c
// 再帰的にトライ木を RDB から読み込む
void HelloTrieType_LoadRecursive(RedisModuleIO *rdb, TrieTypeNode *n) {
    // 終端情報、どの枝に子ノードがあるかの情報を含んだ整数をロード
    uint64_t u = RedisModule_LoadUnsigned(rdb);
    n->terminal = u & 1;

    TrieTypeNode** cursor = n->children;
    while (u >>= 1) {
        if (u & 1) {
            // 子ノードのメモリ割り当てと子ノードを RDB から読み込む
            *cursor = RedisModule_Calloc(1, sizeof(**cursor));
            HelloTrieType_LoadRecursive(rdb, *cursor);
        }
        ++cursor;
    }
}

void *HelloTrieType_Load(RedisModuleIO *rdb, int encver) {
    if (encver != 0) return NULL;

    // ルートノードのメモリ割り当てと RDB から読み込み
    TrieTypeNode *n = RedisModule_Calloc(1, sizeof(*n));
    HelloTrieType_LoadRecursive(rdb, n);
    return n;
}
```

`RedisModule_Load*` 系の関数で、RDB から整数、文字列、浮動小数などを読み込むことが出来るので、後はそれらを使ってデータを復号すればよいです。また、その際のメモリ割り当てには Redis が提供している `RedisModule_Alloc` などの関数を使います。
今回は一つのノードが終端かどうかとどの 26 個の枝があるかという 27bit の情報しか持っていないので、符号なし整数で符号化しています。

RDB からの読み込みが実装できたので、次は RDB への書き込み関数を見てみましょう。

```c:hello_trie.c
void HelloTrieType_Save(RedisModuleIO *rdb, void *value) {
    TrieTypeNode *n = value;

    // 終端情報、子ノード情報を符号化
    uint64_t u = 0;
    TrieTypeNode** cursor = n->children + 26;
    while (cursor-- != n->children) {
        if (*cursor) u |= 1;
        u <<= 1;
    }
    u |= n->terminal;

    // RDB に書き込む
    RedisModule_SaveUnsigned(rdb, u);

    // 子ノードも再帰的に書き込む
    while(u >>= 1) {
        ++cursor;
        if (u & 1) HelloTrieType_Save(rdb, *cursor);
    }
}
```

`RedisModule_Load*` 系の関数と対になる `RedisModule_Save*` 系の関数を使って RDB に書き込むだけです。

続いては、データを AOF に書き出すための関数です。こちらは、書き出したいデータを再構築することのできる Redis のコマンド列を吐き出せばよいです。ここでは、後で定義するトライ木に単語を追加するコマンド `hello.trie.insert` を出力しています。

```c:hello_trie.c
char *HelloTrieType_RewriteRecursive(RedisModuleIO *aof, RedisModuleString *key, TrieTypeNode *n, char *buffer, int depth) {
    // 単語バッファのサイズを増やす
    buffer = RedisModule_Realloc(buffer, sizeof(char) * (depth + 1));

    // 終端だったら単語をトライ木に追加するコマンドを吐く
    if (n->terminal && depth > 0) {
        buffer[depth] = '\0';
        RedisModule_EmitAOF(aof, "hello.trie.insert", "sc", key, buffer);
    }

    // 子ノードを再帰的に処理する
    TrieTypeNode** cursor = n->children;
    char ch = 'a';
    while (cursor != n->children + 26) {
        if (*cursor) {
            buffer[depth] = ch;
            buffer = HelloTrieType_RewriteRecursive(aof, key, *cursor, buffer, depth+1);
        }
        ++cursor;
        ++ch;
    }

    return buffer;
}

void HelloTrieType_Rewrite(RedisModuleIO *aof, RedisModuleString *key, void *value) {
    // 単語用のバッファ領域確保
    char *buffer = RedisModule_Calloc(1, sizeof(char));

    buffer = HelloTrieType_RewriteRecursive(aof, key, value, buffer, 0);

    // 処理が終わったらちゃんとバッファ領域を解放
    RedisModule_Free(buffer);
}
```

トライ木に追加する単語を書き込むための文字列領域を確保して、適宜 `realloc` しつつ処理を進めています。`RedisModule_EmitAOF` がコマンドを書き出す関数で、第 2 引数がコマンド名、第 4 引数以降がコマンドに渡す引数列になっています。第 3 引数の謎の `"sc"` が何を表しているかというと、コマンドに渡す引数列のフォーマットを指定するものです。`s` が `RedisModuleString` で、`c` がヌル文字で終わるいわゆる C の文字列です。他にも整数 `l`、文字列の配列 `v` などが指定できます。

続いては `DEBUG DIGEST` 用のダイジェスト関数なのですが、そもそも `DEBUG DIGEST` がまだモジュールに対応していないようなので今回はスキップで！

```c:hello_trie.c
void HelloTrieType_Digest(RedisModuleDigest *digest, void *value) {
}
```

そして、最後がトライ木のメモリを開放する関数です。

```c:hello_trie.c
void HelloTrieType_Free(void *value) {
    TrieTypeNode *n = value;

    TrieTypeNode** cursor = n->children;
    while (cursor != n->children + 26) {
        if (*cursor) HelloTrieType_Free(*cursor);
        ++cursor;
    }
    RedisModule_Free(n->children);
    RedisModule_Free(n);
}
```

以上で定義した関数を引数として、`RedisModule_CreateDataType` を `RedisModule_OnLoad` で呼び出すことで、晴れて独自データ型を Redis で扱えるようになります！

```c:hello_trie.c
static RedisModuleType *TrieType;

int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
    if (RedisModule_Init(ctx, "hello", 1, REDISMODULE_APIVER_1) == REDISMODULE_ERR)
        return REDISMODULE_ERR;

    TrieType = RedisModule_CreateDataType(ctx, "hellotrie", 0, HelloTrieType_Load, HelloTrieType_Save,
        HelloTrieType_Rewrite, HelloTrieType_Digest, HelloTrieType_Free);
    if (TrieType == NULL)
        return REDISMODULE_ERR;

    // ...
}
```

### トライ木を操作するコマンド

以上で独自のトライ木のデータ型を Redis のメモリ、RDB、AOF で取り扱えるようになったので、ここから具体的にトライ木を操作するためのコマンドを実装していきます。
そろそろ長くなってきたので、今回はトライ木に単語を追加する `hello.trie.insert` と単語のプレフィックスを補完する `hello.trie.complete` の二つだけを作ります。

その前に、トライ木を操作する基本的な関数をいくつか定義します。

```c:hello_trie.c
// トライ木に単語を追加する
void TrieTypeInsert(TrieTypeNode *n, const char *word) {
    while (*word) {
        uint8_t i = *word - 'a';
        if (!n->children[i])
            n->children[i] = RedisModule_Calloc(1, sizeof(TrieTypeNode));
        n = n->children[i];
        ++word;
    }
    n->terminal = 1;
}

// プレフィックスを含む、辞書順で最初に見つかる単語を返す
char *TrieTypeComplete(TrieTypeNode *n, const char *prefix, size_t len, char *result, size_t *newlen) {
    *newlen = 0;
    result = RedisModule_Realloc(result, sizeof(char) * (len + 1));

    while (*prefix) {
        n = n->children[*prefix - 'a'];
        if (!n) return NULL;
        result[(*newlen)++] = *(prefix++);
    }

    while (!n->terminal) {
        result[*newlen] = 'a';

        TrieTypeNode** cursor = n->children;
        while (!*cursor) ++cursor, ++result[*newlen];

        n = *cursor;
        ++*newlen;
        result = RedisModule_Realloc(result, sizeof(char) * (*newlen + 1));
    }
    result[*newlen] = '\0';

    return result;
}
```

これらの処理はトライ木の話であって、Redis あまり関係ないので詳しくは触れません。後は、これらの関数を使った Redis コマンドの本体を実装します。

#### hello.trie.insert

```c:hello_trie.c
// hello.trie.insert KEY WORD
int HelloTrieInsert_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
    if (argc != 3) return RedisModule_WrongArity(ctx);

    RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_READ | REDISMODULE_WRITE);

    // TYPE コマンドで key のタイプを取得
    int type = RedisModule_KeyType(key);

    // キーの型が TrieType じゃなかったらエラー
    if (type != REDISMODULE_KEYTYPE_EMPTY && RedisModule_ModuleTypeGetType(key) != TrieType)
        return RedisModule_ReplyWithError(ctx, REDISMODULE_ERRORMSG_WRONGTYPE);

    TrieTypeNode *n;
    if (type == REDISMODULE_KEYTYPE_EMPTY) {
        // キーが空だったら、空の木を作る
        n = RedisModule_Calloc(1, sizeof(*n));
        RedisModule_ModuleTypeSetValue(key, TrieType, n);
    } else {
        // キーに格納されたトライ木を取得する
        n = RedisModule_ModuleTypeGetValue(key);
    }

    // RedisModuleString から char* に変換
    size_t len;
    const char *word = RedisModule_StringPtrLen(argv[2], &len);

    // トライ木に単語を追加
    TrieTypeInsert(n, word);

    // 返り値はとりあえず NULL
    RedisModule_ReplyWithNull(ctx);
    RedisModule_CloseKey(key);

    // コマンドをスレーブにも伝える
    RedisModule_ReplicateVerbatim(ctx);

    return REDISMODULE_OK;
}
```

ここで `RedisModule_ModuleTypeSetValue` や `RedisModule_ModuleTypeGetValue` を使って、キーに対してトライ木型のデータを格納したり、取得したりしています。また、`RedisModule_ReplicateVerbatim` を呼び出すことで、コマンドをスレーブでも実行させることが出来ます。

#### hello.trie.complete

```c:hello_trie.c
// hello.trie.insert KEY PREFIX
int HelloTrieComplete_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
    if (argc != 3) return RedisModule_WrongArity(ctx);

    RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_READ);

    int type = RedisModule_KeyType(key);
    if (type != REDISMODULE_KEYTYPE_EMPTY && RedisModule_ModuleTypeGetType(key) != TrieType)
        return RedisModule_ReplyWithError(ctx, REDISMODULE_ERRORMSG_WRONGTYPE);

    TrieTypeNode *n;
    if (type == REDISMODULE_KEYTYPE_EMPTY) {
        // 空だったら NULL を返す
        RedisModule_ReplyWithNull(ctx);
    } else {
        n = RedisModule_ModuleTypeGetValue(key);

        size_t len;
        const char *prefix = RedisModule_StringPtrLen(argv[2], &len);

        // 補完結果を格納する領域の確保
        char *result = RedisModule_Alloc(sizeof(char) * (len + 1));

        if (result = TrieTypeComplete(n, prefix, len, result, &len)) {
            // 補完に成功したらその文字列を返す
            RedisModuleString *s = RedisModule_CreateString(ctx, result, len);
            RedisModule_ReplyWithString(ctx, s);
            RedisModule_Free(s);
        } else {
            // 補完できなかったら NULL を返す
            RedisModule_ReplyWithNull(ctx);
        }

        RedisModule_Free(result);
    }

    RedisModule_CloseKey(key);

    return REDISMODULE_OK;
}
```

こちらは今まで見てきた関数を使ってるだけなので特筆すべきことはないです。

これで、コマンドの本体を実装できたので、後はこれをコマンド名と紐付けるだけです。

```c:hello_trie.c
int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
    // ...

    if (RedisModule_CreateCommand(ctx, "hello.trie.insert",
            HelloTrieInsert_RedisCommand, "write deny-oom", 1, 1, 1) == REDISMODULE_ERR)
        return REDISMODULE_ERR;
  
    if (RedisModule_CreateCommand(ctx, "hello.trie.complete",
            HelloTrieComplete_RedisCommand, "readonly", 1, 1, 1) == REDISMODULE_ERR)
        return REDISMODULE_ERR;

    return REDISMODULE_OK;
}
```

### 動作確認

```redis
127.0.0.1:6379> hello.trie.insert trie hello
(nil)
127.0.0.1:6379> hello.trie.insert trie helloworld
(nil)
127.0.0.1:6379> hello.trie.complete trie he
"hello"
127.0.0.1:6379> hello.trie.complete trie helo
(nil)
127.0.0.1:6379> hello.trie.complete trie hello
"hello"
127.0.0.1:6379> hello.trie.complete trie hellow
"helloworld"
127.0.0.1:6379> hello.trie.complete trie world
(nil)
```

ちゃんと動いてる！RDB、AOF、レプリケーション周りも全部うまく行っている(完)！

## まとめ

今後どうなっていくのか全然分からないですが、Redis モジュール書くの楽しかったです！
今なら開発してる人も少ないそうだし、将来メジャーになりうるモジュールを作るいいチャンスかもしれないです。

ちなみに今回実装したトライ木のソースコードは [hello/trie.c](https://github.com/saidie/redis_modules/blob/b96ffdfef96f3a0067149b762194ec1be88db076/hello/trie.c) にあります (コマンドが多かったり微妙に違いますが)。