[アップデート]RDS for PostgreSQLがplrust拡張機能に対応しました

初めに

先日のアップデートでRDS for PostgreSQLで利用可能な拡張機能にplrustが追加されました。

この拡張機能を利用することで関数の作成にRustを利用することが可能となります。

PostgreSQL 15.2以上のみの対応なります。

信頼できるRust

さてRustは言語仕様上ファイルシステムへのアクセスなどAWS側で管理すべき部分にアクセスが可能であり、またunsafeブロックを利用することでシステムを触るような危険なコードを作成することが可能です。

この部分の処理からどのように保護されているかが気になります。

https://tcdi.github.io/plrust/trusted-untrusted.html
Normally, PL/Rust is installed as a "trusted" programming language named plrust. In this setup, certain Rust and pgrx operations are disabled to preserve security. In general, the operations that are restricted are those that interact with the environment. This includes file handle operations, require, and use (for external modules). There is no way to access internals of the database server process or to gain OS-level access with the permissions of the server process, as a C function can do. Thus, any unprivileged database user can be permitted to use this language.

ドキュメントを読んでみたところplrustの機能としてplrustのコンパイル時にコンパイルターゲットとして"trusted"か"untrasted"かどちらのRustを利用するか選択ができます。

"trusted"なplrustを選択しコンパイルする場合はファイルシステムやOS情報の取得などのデータベースサーバ自体の情報取得やネットワーク通信など外部通信が制限されます。
RDSのドキュメント上で明言はされていませんがPL/Rust is a trusted Rust language extension for PostgreSQL.と書いてあるので指定時はこちらが利用されているのではないかと想像しています。

https://tcdi.github.io/plrust/plrust.html#what-about-unsafe
PL/Rust uses the Rust compiler itself to wholesale disallow the use of unsafe in user functions. If a LANGUAGE plrust function uses unsafe it won't compile.

なおunsafeについてはオプションによらず禁止されてるようです。
厳密にはサードパーティーのcrateにunsafeが含まれることは許可されるようですが仕組み的に機能が無効化されるようです。

実行環境

• RDS for PostgreSQL 15.3-R

postgres=> SELECT version();
                                                 version
---------------------------------------------------------------------------------------------------------
 PostgreSQL 15.3 on x86_64-pc-linux-gnu, compiled by gcc (GCC) 7.3.1 20180712 (Red Hat 7.3.1-12), 64-bit
(1 row)

インストール

デフォルトのパラメータグループではshared_preload_librariesにplrustが追加されていないためパラメータの追加を行い、インスタンスの再起動を行う必要があります。

再起動後CREATE EXTENSIONでplrustを追加します。

postgres=> CREATE EXTENSION plrust;
CREATE EXTENSION
postgres=> select * from pg_extension ;
  oid  | extname | extowner | extnamespace | extrelocatable | extversion | extconfig | extcondition
-------+---------+----------+--------------+----------------+------------+-----------+--------------
 14498 | plpgsql |       10 |           11 | f              | 1.0        |           |
 16442 | plrust  |       10 |        16441 | f              | 1.0        |           |
(2 rows)

作成と実行

シンプルに１から受け取った引数までの整数の総和を求めるsum_1_to_x(x)を作成し実行してみます。

postgres=> \timing
Timing is on.
postgres=> CREATE OR REPLACE FUNCTION sum_1_to_x(x int)
RETURNS int
LANGUAGE plrust
IMMUTABLE PARALLEL SAFE STRICT
AS $$
    Ok(Some((1..=x).sum()))
$$;
CREATE FUNCTION
Time: 11017.555 ms (00:11.018)
postgres=> SELECT sum_1_to_x(10);
 sum_1_to_x
------------
         55
(1 row)

コンパイルはCREATE FUNCITONの実行時に行われるためrustのコードはそのまま記載して問題ありません。

ただコンパイルでCPU利用率やメモリ使用率への影響が考えられるため実行タイミングには気をつけた方が良いかもしれません。

ファイルシステムへのアクセスを試みる

せっかくなので試してみたいですよね？

現在のディレクトリを表示するpwd()関数を作成し実行します。

postgres=> CREATE OR REPLACE FUNCTION pwd()
postgres-> RETURNS TEXT
postgres-> LANGUAGE plrust
postgres-> IMMUTABLE PARALLEL SAFE STRICT
postgres-> AS $$
postgres$>     Ok(Some(std::env::current_dir()?.to_string_lossy().to_string()))
postgres$> $$;
CREATE FUNCTION
Time: 9462.428 ms (00:09.462)
postgres=> SELECT pwd();
ERROR:  operation not supported on this platformTime: 24.688 ms

どうやらコンパイル時点でエラーとなるものではなく、実際の操作時にエラーとなるようです。

なおunsafeブロックが含まれる場合はドキュメント記載の通りコンパイル自体が失敗します。

postgres=> CREATE OR REPLACE FUNCTION pwd()
postgres-> RETURNS TEXT
postgres-> LANGUAGE plrust
postgres-> IMMUTABLE PARALLEL SAFE STRICT
postgres-> AS $$
postgres$>     unsafe {
postgres$>         Ok(Some(std::env::current_dir()?.to_string_lossy().to_string()))
postgres$>     }
postgres$> $$;
ERROR:
   0: `cargo build` failed
...
`cargo build` stderr:
      Compiling plrust_fn_oid_5_24649_6781753360384 v0.0.0 (/rdsdbdata/extensions/plrust/plrust_fn_oid_5_24649_6781753360384)
   error: usage of an `unsafe` block
     --> src/lib.rs:43:9
      |
   43 |         unsafe { Ok(Some(std::env::current_dir()?.to_string_lossy().to_string())) }
      |         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
      |
   note: the lint level is defined here
     --> src/lib.rs:34:15
      |
   34 |     #![forbid(unsafe_code)]
      |               ^^^^^^^^^^^

返却型値について

plrustで作成した関数の返却は以下の型で行う必要があります。

Result<Option<T>,　Box<dyn std::error::Error + Send + Sync + 'static>>

Rustにはnullがない関係でPostgreSQLにnullを渡したい場合は代替としてNoneを渡していることで実現しており、そのためにOption型で指定する必要があるようです。

ResultについてはRustは例外処理が存在しないのでErrを明示的に呼び出すことで処理をエラーとして上位に伝達し呼び出し元でのトランザクションの中止等を行います。

ただ今回の検証の中でpanic!()を引き起こした場合とErr()を返却した場合の違いについては追いきれませんでした（ソースコード少し追って一ましたが定義場所が追いきれず...)。

先ほどのコードで?オペレータの代わりにunwrap()でpanicを引き起こすようにしましたが、メッセージ上の違いはあるものの同様にトランザクションは中断され明確に機能的にここが気になるような点は見当たりませんでした。

$ SELECT pwd();
ERROR:  called `Result::unwrap()` on an `Err` value: Error { kind: Unsupported, message: "operation not supported on this platform" }

ドキュメントを読む限りErr()の返却を想定しているため、Err()を返却しておく方が無難そうです。

速度

AWS Blogの例ではJavaScriptに比べRustの処理が約17倍早くなるようです。

今回はdb.t3.microでCPUクレジットが問題にならないように十分残っている状態で確認しております。

素数計算

同じ方法で試すのも寂しいのですぐ思いついた方法で素数算出をやってもらいました。
(コードは差異が出にくいようにChatGPTに作成してもらっています)

CREATE OR REPLACE FUNCTION get_nth_prime_rust(n int)
RETURNS int
LANGUAGE plrust
IMMUTABLE PARALLEL SAFE STRICT
AS $$
    let mut primes = Vec::new();
    let mut num = 2;
    
    let is_prime = |x: i32| -> bool {
        if x < 2 {
            return false;
        }
        
        for i in 2..=(x as f64).sqrt() as i32 {
            if x % i == 0 {
                return false;
            }
        }
        
        true
    };
    
    while primes.len() < n as usize {
        if is_prime(num) {
            primes.push(num);
        }
        num += 1;
    }
    
    Ok(Some(primes[n as usize - 1]))
$$;

CREATE OR REPLACE FUNCTION get_nth_prime_js(n int)
RETURNS int
LANGUAGE plv8
IMMUTABLE PARALLEL SAFE STRICT
AS $$
  const primes = [];
  let num = 2;

  const isPrime = (x) => {
    if (x < 2) {
      return false;
    }

    for (let i = 2; i <= Math.sqrt(x); i++) {
      if (x % i === 0) {
        return false;
      }
    }

    return true;
  };

  while (primes.length < n) {
    if (isPrime(num)) {
      primes.push(num);
    }
    num++;
  }

  return primes[n - 1];
$$;

何度か実行してキャッシュが聞いているような状態だとこのケースではほとんど差はできないようです。

postgres=> EXPLAIN ANALYZE SELECT get_nth_prime_rust(500000);
                                     QUERY PLAN
------------------------------------------------------------------------------------
 Result  (cost=0.00..0.01 rows=1 width=4) (actual time=0.001..0.002 rows=1 loops=1)
 Planning Time: 6019.998 ms
 Execution Time: 0.021 ms
(3 rows)
postgres=> EXPLAIN ANALYZE SELECT get_nth_prime_js(500000);
                                     QUERY PLAN
------------------------------------------------------------------------------------
 Result  (cost=0.00..0.01 rows=1 width=4) (actual time=0.001..0.002 rows=1 loops=1)
 Planning Time: 6198.226 ms
 Execution Time: 0.018 ms
(3 rows)

行列計算

行列を計算してもらったところこちらは有意な差が見られました。

CREATE OR REPLACE FUNCTION calculate_matrix_rust(matrix_size int)
RETURNS int
LANGUAGE plrust
IMMUTABLE PARALLEL SAFE STRICT
AS $$
    // 行列の初期化
    let mut matrix: Vec<Vec<i32>> = vec![vec![0; matrix_size as usize]; matrix_size as usize];

    // 行列の計算
    for i in 0..matrix_size {
        for j in 0..matrix_size {
            matrix[i as usize][j as usize] = (i * j) as i32;
        }
    }

    // 行列の掛け算
    let mut result: Vec<Vec<i32>> = vec![vec![0; matrix_size as usize]; matrix_size as usize];
    for i in 0..matrix_size {
        for j in 0..matrix_size {
            for k in 0..matrix_size {
                result[i as usize][j as usize] += matrix[i as usize][k as usize] * matrix[k as usize][j as usize];
            }
        }
    }

    // 最後の要素の値を返す
    Ok(Some(result[(matrix_size - 1) as usize][(matrix_size - 1) as usize]))
$$;

CREATE OR REPLACE FUNCTION calculate_matrix_js(matrix_size int)
RETURNS int
LANGUAGE plv8
IMMUTABLE PARALLEL SAFE STRICT
AS $$
  // 行列の初期化
  const matrix = new Array(matrix_size);
  for (let i = 0; i < matrix_size; i++) {
    matrix[i] = new Array(matrix_size).fill(0);
  }

  // 行列の計算
  for (let i = 0; i < matrix_size; i++) {
    for (let j = 0; j < matrix_size; j++) {
      matrix[i][j] = i * j;
    }
  }

  // 行列の掛け算
  const result = new Array(matrix_size);
  for (let i = 0; i < matrix_size; i++) {
    result[i] = new Array(matrix_size).fill(0);
    for (let j = 0; j < matrix_size; j++) {
      for (let k = 0; k < matrix_size; k++) {
        result[i][j] += matrix[i][k] * matrix[k][j];
      }
    }
  }

  // 最後の要素の値を返す
  return result[matrix_size - 1][matrix_size - 1];
$$;

postgres=> EXPLAIN ANALYZE SELECT * FROM calculate_matrix_js(1000);
                                     QUERY PLAN
------------------------------------------------------------------------------------
 Result  (cost=0.00..0.01 rows=1 width=4) (actual time=0.001..0.001 rows=1 loops=1)
 Planning Time: 15122.189 ms
 Execution Time: 0.018 ms
(3 rows)

postgres=> EXPLAIN ANALYZE SELECT * FROM calculate_matrix_rust(1000);
                                     QUERY PLAN
------------------------------------------------------------------------------------
 Result  (cost=0.00..0.01 rows=1 width=4) (actual time=0.001..0.002 rows=1 loops=1)
 Planning Time: 6698.201 ms
 Execution Time: 0.018 ms
(3 rows)

終わりに

今回はRDS for PostgreSQLにplrustが対応したため試してみました。

plrustは存在は知ってはいたものの環境を整えるのが手間で触っていなかったのですがRDSであれば設定値を変えるだけで手軽に導入できるので、せっかくなので触ってみようができる良い機会でした。

当ブログではAWS Lambdaのを壊しに行こうとする方もいらっしゃいますが、残念ながら(?)plrustではファイルシステムへのアクセス等がきちんと制限される仕組みが組み込まれているため今回は難しそうです。

つまりのところしっかりマネージドサービスとして責任範囲の切り分けとしてさわれない部分をしっかり分けてくれているのは安心点ではあります。

また、計算速度については普段ベンチマークのようなことはやらないので勝手なイメージですが、素数算出みたいな単純計算はRustの方が早そうな印象があったのでJavaScriptと同等というのは意外でした。

この辺りについてはPostgreSQLとしての最適化部分もあるかと思いますのでもっと色々なケースを試してどのようなケースで差異が出やすいかというのは個人的は気になる点です。