Reduce SQL sanitizer allocations #2136

ninedraft · 2024-10-01T14:43:53Z

#2124

Result:

Main optimizations:

extensive usage of sync.Pool for byte buffers, lexers and parsed query structs
append-style string formatters for int64, float64 and time.Time + bytes.Buffer.AvailableBuffer
rework of QuoteString and QuoteBytes to append-style (with tests for backwards compatibility)

Misc changes:

benchmarks for Query.Sanitize and SanitizeSQL functions
a tiny script for generation of benchmark reports for selected commits and diff (using benchstat)
fuzzing of QuoteString and QuoteBytes (I did'n find any problems for 1h of fuzzing, but you can never be sure for 100%)

Since optimization is an extremely hard problem, I think it's worth checking some more benchmarks.

I would be very grateful for your opinion on this and recommendations/advice, @jackc @vtolstov

make benchmark more extensive add quote to string add BenchmarkSanitizeSQL

fix benchmmark script fix benchmark script

check new quoteBytes

use lexer pool

ninedraft · 2024-10-01T15:01:23Z

benchmark diffs for concrete optimisations

goos: darwin
goarch: arm64
pkg: github.com/jackc/pgx/v5/internal/sanitize
cpu: Apple M1
              │ benchmarks/0_base_case.bench │     benchmarks/1_buf_pool.bench     │ benchmarks/2_append_AvailableBuffer.bench │    benchmarks/3_quoteBytes.bench    │   benchmarks/4_quoteString.bench    │ benchmarks/5_add_lexer_and_query_pools.bench │ benchmarks/6_drop_too_large_values_from_memory_pools.bench │
              │            sec/op            │   sec/op     vs base                │      sec/op        vs base                │   sec/op     vs base                │   sec/op     vs base                │        sec/op         vs base                │               sec/op                vs base                │
Sanitize-8                       718.2n ± 1%   578.8n ± 1%  -19.41% (p=0.000 n=10)         439.9n ± 0%  -38.74% (p=0.000 n=10)   413.6n ± 4%  -42.42% (p=0.000 n=10)   397.1n ± 1%  -44.72% (p=0.000 n=10)            403.6n ± 1%  -43.81% (p=0.000 n=10)                          400.8n ± 2%  -44.20% (p=0.000 n=10)
SanitizeSQL-8                    2.089µ ± 0%   1.956µ ± 0%   -6.37% (p=0.000 n=10)         1.828µ ± 0%  -12.49% (p=0.000 n=10)   1.812µ ± 1%  -13.28% (p=0.000 n=10)   1.789µ ± 1%  -14.36% (p=0.000 n=10)            1.670µ ± 0%  -20.06% (p=0.000 n=10)                          1.673µ ± 0%  -19.91% (p=0.000 n=10)
geomean                          1.225µ        1.064µ       -13.13%                        896.8n       -26.79%                  865.5n       -29.34%                  842.8n       -31.19%                           820.9n       -32.98%                                         818.8n       -33.15%

              │ benchmarks/0_base_case.bench │     benchmarks/1_buf_pool.bench     │ benchmarks/2_append_AvailableBuffer.bench │    benchmarks/3_quoteBytes.bench    │   benchmarks/4_quoteString.bench    │ benchmarks/5_add_lexer_and_query_pools.bench │ benchmarks/6_drop_too_large_values_from_memory_pools.bench │
              │             B/op             │    B/op      vs base                │       B/op         vs base                │    B/op      vs base                │    B/op      vs base                │         B/op          vs base                │                B/op                 vs base                │
Sanitize-8                       1488.0 ± 0%    528.0 ± 0%  -64.52% (p=0.000 n=10)          472.0 ± 0%  -68.28% (p=0.000 n=10)    456.0 ± 0%  -69.35% (p=0.000 n=10)    424.0 ± 0%  -71.51% (p=0.000 n=10)             424.0 ± 0%  -71.51% (p=0.000 n=10)                           424.0 ± 0%  -71.51% (p=0.000 n=10)
SanitizeSQL-8                    2216.0 ± 0%   1256.0 ± 0%  -43.32% (p=0.000 n=10)         1200.0 ± 0%  -45.85% (p=0.000 n=10)   1184.0 ± 0%  -46.57% (p=0.000 n=10)   1152.0 ± 0%  -48.01% (p=0.000 n=10)             552.0 ± 0%  -75.09% (p=0.000 n=10)                           552.0 ± 0%  -75.09% (p=0.000 n=10)
geomean                         1.773Ki         814.4       -55.15%                         752.6       -58.55%                   734.8       -59.54%                   698.9       -61.51%                            483.8       -73.36%                                          483.8       -73.36%

              │ benchmarks/0_base_case.bench │    benchmarks/1_buf_pool.bench     │ benchmarks/2_append_AvailableBuffer.bench │   benchmarks/3_quoteBytes.bench    │   benchmarks/4_quoteString.bench   │ benchmarks/5_add_lexer_and_query_pools.bench │ benchmarks/6_drop_too_large_values_from_memory_pools.bench │
              │          allocs/op           │ allocs/op   vs base                │     allocs/op      vs base                │ allocs/op   vs base                │ allocs/op   vs base                │      allocs/op        vs base                │             allocs/op               vs base                │
Sanitize-8                       11.000 ± 0%   7.000 ± 0%  -36.36% (p=0.000 n=10)          4.000 ± 0%  -63.64% (p=0.000 n=10)   3.000 ± 0%  -72.73% (p=0.000 n=10)   2.000 ± 0%  -81.82% (p=0.000 n=10)             2.000 ± 0%  -81.82% (p=0.000 n=10)                           2.000 ± 0%  -81.82% (p=0.000 n=10)
SanitizeSQL-8                     26.00 ± 0%   22.00 ± 0%  -15.38% (p=0.000 n=10)          19.00 ± 0%  -26.92% (p=0.000 n=10)   18.00 ± 0%  -30.77% (p=0.000 n=10)   17.00 ± 0%  -34.62% (p=0.000 n=10)             10.00 ± 0%  -61.54% (p=0.000 n=10)                           10.00 ± 0%  -61.54% (p=0.000 n=10)
geomean                           16.91        12.41       -26.62%                         8.718       -48.45%                  7.348       -56.55%                  5.831       -65.52%                            4.472       -73.56%                                          4.472       -73.56%

jackc · 2024-10-05T14:50:38Z

LGTM. But this is obviously a very security critical part of the code, so I'd like if we can get some more eyes on this before merging.

vtolstov · 2024-10-06T13:05:58Z

lgtm, i'm try to check on my hot path in next few days.

vtolstov · 2024-10-10T21:05:00Z

In my tests i don't saw any issues.

ninedraft · 2024-10-15T17:03:52Z

@jackc

But this is obviously a very security critical part of the code, so I'd like if we can get some more eyes on this before merging.

It would be very much appreciated if you could suggest someone I can tag on this issue. I'm also in the process of writing more tests for SQL injection + more fuzzing

jackc · 2024-10-18T21:45:56Z

It would be very much appreciated if you could suggest someone I can tag on this issue.

I wish I could. Unfortunately, I don't know of anyone.

I'm also in the process of writing more tests for SQL injection + more fuzzing

👍 It's been a couple weeks since I reviewed the code, so I can review it again with fresh eyes now. It's not quite as good as multiple reviewers, but at least it will be multiple reviews.

I'll wait until you add the additional tests.

sean-

See the inline suggested changes. These small set of optimizations reduced the sec/op from 607.9n down to 604.9n (-0.49%). I can push this as a different PR or these changes can be incorporated into this branch.

$ bash benchmmark.sh 2ec900454bfe65daa9648488e93f7627c26b810c 82642726914a8b054ca123fd87c4d984da6d78eb 431e11b61c809c2373128ecf63ed48cf8bdf4dd4 71c3b107187b02ea44dbc7d38e931115ca7286c7
$ benchstat benchmarks/*.bench
goos: darwin
goarch: arm64
pkg: github.com/jackc/pgx/v5/internal/sanitize
cpu: Apple M3 Pro
               │ benchmarks/1_fix_preallocations_of_quoted_string.bench │ benchmarks/2_optimize_QuoteBytes.bench │ benchmarks/2_rework_QuoteString_to_iterate_over_bytes.bench │ benchmarks/3_optimize_QuoteBytes.bench │
               │                         sec/op                         │     sec/op       vs base               │                sec/op                 vs base               │     sec/op       vs base               │
Sanitize-12                                                 307.1n ± 1%       300.6n ± 2%  -2.10% (p=0.001 n=10)                            305.3n ± 1%  -0.57% (p=0.015 n=10)       304.2n ± 1%  -0.93% (p=0.003 n=10)
SanitizeSQL-12                                              1.204µ ± 2%       1.207µ ± 1%       ~ (p=0.100 n=10)                            1.204µ ± 2%       ~ (p=0.697 n=10)       1.203µ ± 2%       ~ (p=0.898 n=10)
geomean                                                     607.9n            602.3n       -0.91%                                           606.3n       -0.26%                      604.9n       -0.49%

               │ benchmarks/1_fix_preallocations_of_quoted_string.bench │ benchmarks/2_optimize_QuoteBytes.bench │ benchmarks/2_rework_QuoteString_to_iterate_over_bytes.bench │ benchmarks/3_optimize_QuoteBytes.bench │
               │                          B/op                          │     B/op       vs base                 │                B/op                 vs base                 │     B/op       vs base                 │
Sanitize-12                                                  424.0 ± 0%      424.0 ± 0%       ~ (p=1.000 n=10) ¹                           424.0 ± 0%       ~ (p=1.000 n=10) ¹      424.0 ± 0%       ~ (p=1.000 n=10) ¹
SanitizeSQL-12                                               552.0 ± 0%      552.0 ± 0%       ~ (p=1.000 n=10) ¹                           552.0 ± 0%       ~ (p=1.000 n=10) ¹      552.0 ± 0%       ~ (p=1.000 n=10) ¹
geomean                                                      483.8           483.8       +0.00%                                            483.8       +0.00%                       483.8       +0.00%
¹ all samples are equal

               │ benchmarks/1_fix_preallocations_of_quoted_string.bench │ benchmarks/2_optimize_QuoteBytes.bench │ benchmarks/2_rework_QuoteString_to_iterate_over_bytes.bench │ benchmarks/3_optimize_QuoteBytes.bench │
               │                       allocs/op                        │   allocs/op    vs base                 │             allocs/op               vs base                 │   allocs/op    vs base                 │
Sanitize-12                                                  2.000 ± 0%      2.000 ± 0%       ~ (p=1.000 n=10) ¹                           2.000 ± 0%       ~ (p=1.000 n=10) ¹      2.000 ± 0%       ~ (p=1.000 n=10) ¹
SanitizeSQL-12                                               10.00 ± 0%      10.00 ± 0%       ~ (p=1.000 n=10) ¹                           10.00 ± 0%       ~ (p=1.000 n=10) ¹      10.00 ± 0%       ~ (p=1.000 n=10) ¹
geomean                                                      4.472           4.472       +0.00%                                            4.472       +0.00%                       4.472       +0.00%
¹ all samples are equal

sean- · 2024-10-21T15:22:53Z

internal/sanitize/benchmmark.sh

+    }
+
+    # Sanitized commmit message
+    commit_message=$(git log -1 --pretty=format:"%s" | tr ' ' '_')


This needs to escape /:

commit_message=$(git log -1 --pretty=format:"%s" | tr -c '[:alnum:]-_' '_')

sean- · 2024-10-21T15:23:51Z

internal/sanitize/benchmmark.sh

+    bench_files+=("$bench_file")
+done
+
+benchstat "${bench_files[@]}"


Can you prefix with a small comment: # go install golang.org/x/perf/cmd/benchstat@latest

sean- · 2024-10-21T15:26:50Z

internal/sanitize/sanitize.go

+
+	dst = append(dst, quote...)
+
+	return dst
 }


This is purely a style nit, but I don't like reslicing for these types of functions because it's not idiomatic and hard to follow. I took the above QuoteString() and replaced it with something that uses an iterator:

func QuoteString(dst []byte, str string) []byte { const quote = '\'' // Preallocate space for the worst case scenario dst = slices.Grow(dst, len(str)*2+2) // Add opening quote dst = append(dst, quote) // Iterate through the string without allocating for i := 0; i < len(str); i++ { if str[i] == quote { dst = append(dst, quote, quote) } else { dst = append(dst, str[i]) } } // Add closing quote dst = append(dst, quote) return dst }

sean- · 2024-10-21T15:27:39Z

internal/sanitize/sanitize.go

+	dst = append(dst, p...)
+
+	dst = append(dst, `'`...)
+	return dst
 }


I was able to measure an improvement by optimizing this function:

func QuoteBytes(dst, buf []byte) []byte { if len(buf) == 0 { return append(dst, `'\x'`...) } // Calculate required length requiredLen := 3 + hex.EncodedLen(len(buf)) + 1 // Ensure dst has enough capacity if cap(dst)-len(dst) < requiredLen { newDst := make([]byte, len(dst), len(dst)+requiredLen) copy(newDst, dst) dst = newDst } // Record original length and extend slice origLen := len(dst) dst = dst[:origLen+requiredLen] // Add prefix dst[origLen] = '\'' dst[origLen+1] = '\\' dst[origLen+2] = 'x' // Encode bytes directly into dst hex.Encode(dst[origLen+3:len(dst)-1], buf) // Add suffix dst[len(dst)-1] = '\'' return dst }

sean-

See the inline suggested changes. These small set of optimizations reduced the sec/op from 607.9n down to 604.9n (-0.49%). I can push this as a different PR or these changes can be incorporated into this branch.

$ bash benchmmark.sh 2ec900454bfe65daa9648488e93f7627c26b810c 82642726914a8b054ca123fd87c4d984da6d78eb 431e11b61c809c2373128ecf63ed48cf8bdf4dd4 71c3b107187b02ea44dbc7d38e931115ca7286c7
$ benchstat benchmarks/*.bench
goos: darwin
goarch: arm64
pkg: github.com/jackc/pgx/v5/internal/sanitize
cpu: Apple M3 Pro
               │ benchmarks/1_fix_preallocations_of_quoted_string.bench │ benchmarks/2_optimize_QuoteBytes.bench │ benchmarks/2_rework_QuoteString_to_iterate_over_bytes.bench │ benchmarks/3_optimize_QuoteBytes.bench │
               │                         sec/op                         │     sec/op       vs base               │                sec/op                 vs base               │     sec/op       vs base               │
Sanitize-12                                                 307.1n ± 1%       300.6n ± 2%  -2.10% (p=0.001 n=10)                            305.3n ± 1%  -0.57% (p=0.015 n=10)       304.2n ± 1%  -0.93% (p=0.003 n=10)
SanitizeSQL-12                                              1.204µ ± 2%       1.207µ ± 1%       ~ (p=0.100 n=10)                            1.204µ ± 2%       ~ (p=0.697 n=10)       1.203µ ± 2%       ~ (p=0.898 n=10)
geomean                                                     607.9n            602.3n       -0.91%                                           606.3n       -0.26%                      604.9n       -0.49%

               │ benchmarks/1_fix_preallocations_of_quoted_string.bench │ benchmarks/2_optimize_QuoteBytes.bench │ benchmarks/2_rework_QuoteString_to_iterate_over_bytes.bench │ benchmarks/3_optimize_QuoteBytes.bench │
               │                          B/op                          │     B/op       vs base                 │                B/op                 vs base                 │     B/op       vs base                 │
Sanitize-12                                                  424.0 ± 0%      424.0 ± 0%       ~ (p=1.000 n=10) ¹                           424.0 ± 0%       ~ (p=1.000 n=10) ¹      424.0 ± 0%       ~ (p=1.000 n=10) ¹
SanitizeSQL-12                                               552.0 ± 0%      552.0 ± 0%       ~ (p=1.000 n=10) ¹                           552.0 ± 0%       ~ (p=1.000 n=10) ¹      552.0 ± 0%       ~ (p=1.000 n=10) ¹
geomean                                                      483.8           483.8       +0.00%                                            483.8       +0.00%                       483.8       +0.00%
¹ all samples are equal

               │ benchmarks/1_fix_preallocations_of_quoted_string.bench │ benchmarks/2_optimize_QuoteBytes.bench │ benchmarks/2_rework_QuoteString_to_iterate_over_bytes.bench │ benchmarks/3_optimize_QuoteBytes.bench │
               │                       allocs/op                        │   allocs/op    vs base                 │             allocs/op               vs base                 │   allocs/op    vs base                 │
Sanitize-12                                                  2.000 ± 0%      2.000 ± 0%       ~ (p=1.000 n=10) ¹                           2.000 ± 0%       ~ (p=1.000 n=10) ¹      2.000 ± 0%       ~ (p=1.000 n=10) ¹
SanitizeSQL-12                                               10.00 ± 0%      10.00 ± 0%       ~ (p=1.000 n=10) ¹                           10.00 ± 0%       ~ (p=1.000 n=10) ¹      10.00 ± 0%       ~ (p=1.000 n=10) ¹
geomean                                                      4.472           4.472       +0.00%                                            4.472       +0.00%                       4.472       +0.00%
¹ all samples are equal

sean-

~~[review comment was posted twice for some reason]~~

ninedraft added 14 commits October 1, 2024 16:54

base case

b9d0214

make benchmark more extensive add quote to string add BenchmarkSanitizeSQL

add benchmark tool

730c324

fix benchmmark script fix benchmark script

buf pool

c57e0d6

shared bytestring

22e4205

append AvailableBuffer

9639283

docs

f5c9af0

quoteBytes

4dcba02

check new quoteBytes

quoteString

cc0b941

decrease number of samples in go benchmark

3f27c12

add FuzzQuoteString and FuzzQuoteBytes

01e234e

add lexer and query pools

07809d5

use lexer pool

rework QuoteString and QuoteBytes as append-style

48cc36a

add docs to sanitize tests

1d50b82

drop too large values from memory pools

339b193

ninedraft marked this pull request as ready for review October 1, 2024 14:52

ninedraft added 2 commits October 20, 2024 18:00

add prefix to quoters tests

e5053ee

fix preallocations of quoted string

8264272

sean- reviewed Oct 21, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Reduce SQL sanitizer allocations #2136

Reduce SQL sanitizer allocations #2136

ninedraft commented Oct 1, 2024

ninedraft commented Oct 1, 2024

jackc commented Oct 5, 2024

vtolstov commented Oct 6, 2024

vtolstov commented Oct 10, 2024

ninedraft commented Oct 15, 2024

jackc commented Oct 18, 2024

sean- left a comment

sean- Oct 21, 2024 •

edited

Loading

sean- Oct 21, 2024

sean- Oct 21, 2024

sean- Oct 21, 2024 •

edited

Loading

sean- left a comment

sean- left a comment •

edited

Loading

Reduce SQL sanitizer allocations #2136

Are you sure you want to change the base?

Reduce SQL sanitizer allocations #2136

Conversation

ninedraft commented Oct 1, 2024

ninedraft commented Oct 1, 2024

jackc commented Oct 5, 2024

vtolstov commented Oct 6, 2024

vtolstov commented Oct 10, 2024

ninedraft commented Oct 15, 2024

jackc commented Oct 18, 2024

sean- left a comment

Choose a reason for hiding this comment

sean- Oct 21, 2024 • edited Loading

Choose a reason for hiding this comment

sean- Oct 21, 2024

Choose a reason for hiding this comment

sean- Oct 21, 2024

Choose a reason for hiding this comment

sean- Oct 21, 2024 • edited Loading

Choose a reason for hiding this comment

sean- left a comment

Choose a reason for hiding this comment

sean- left a comment • edited Loading

Choose a reason for hiding this comment

sean- Oct 21, 2024 •

edited

Loading

sean- Oct 21, 2024 •

edited

Loading

sean- left a comment •

edited

Loading