wrk拡張によるTCPバイトストリームプロトコル対応

wrkはHTTP(s)負荷テストツールとして広く利用されていますが、TCPベースのバイトストリームプロトコルに対応するためにはカスタマイズが必要です。本記事ではその実装方法を解説します。

基本アーキテクチャ

wrkは以下の主要コンポーネントで構成されています：

• Redisのaeイベントループフレームワーク
• LuaJITスクリプティングエンジン
• OpenSSLライブラリ
• HTTPパーサー

Luaスクリプトによる拡張

TCPバイトストリーム対応にはstream_response関数を追加します：

-- TCPレスポンスを処理する関数
-- 例: {"status":0,"message":"success"}
function stream_response(data)
    local response = json.parse(data)
    return response.status == 0
end

パフォーマンス最適化

wrkのマルチスレッドアーキテクチャはRedisのaeフレームワークを利用し、以下の特徴があります：

• スレッドあたり接続数/スレッド数の接続を確立
• イベントループに登録されたfdの監視
• readable/writableイベントの効率的な処理

統計処理の実装

wrkは以下の統計情報を収集します：

typedef struct {
    uint64_t count;     // サンプル数
    uint64_t limit;     // 最大サンプル値
    uint64_t min;       // 最小サンプル値
    uint64_t max;       // 最大サンプル値
    uint64_t data[];    // 値の出現頻度
} stats;

TCPプロトコル対応の実装

主要な変更点はsocket_readable関数内のレスポンス処理です：

static void socket_readable(aeEventLoop *loop, int fd, void *data, int mask) {
    connection *conn = data;
    size_t bytes_read;
    
    do {
        if (!process_tcp_response(conn, &bytes_read))
            goto error_handling;
    } while (bytes_read == BUFFER_SIZE && has_more_data(conn));
    
    return;
    
error_handling:
    conn->thread->errors.read++;
    reconnect(conn);
}

追加機能

拡張版では以下の機能を追加しています：

• JSONパーサー（yyjsonベース）
• MD5ハッシュ生成（nginx/md5ベース）
• TCPテキストプロトコル対応

今後の課題

• 複合テーブルのコピー問題の修正
• TCP接続確立後の認証バイトストリーム送信機能
• UNIXドメインソケット対応