Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Reduce SQL sanitizer allocations #2136

Open
wants to merge 16 commits into
base: master
Choose a base branch
from

Conversation

ninedraft
Copy link
Contributor

#2124

Result:
Pasted Graphic 1

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

@ninedraft ninedraft marked this pull request as ready for review October 1, 2024 14:52
@ninedraft
Copy link
Contributor Author

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
Copy link
Owner

jackc commented Oct 5, 2024

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
Copy link

vtolstov commented Oct 6, 2024

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

@vtolstov
Copy link

In my tests i don't saw any issues.

@ninedraft
Copy link
Contributor Author

@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
Copy link
Owner

jackc commented Oct 18, 2024

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.

Copy link
Contributor

@sean- sean- left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

}

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

@sean- sean- Oct 21, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This needs to escape /:

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

bench_files+=("$bench_file")
done

benchstat "${bench_files[@]}"
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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


dst = append(dst, quote...)

return dst
}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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
}

dst = append(dst, p...)

dst = append(dst, `'`...)
return dst
}
Copy link
Contributor

@sean- sean- Oct 21, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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
}

Copy link
Contributor

@sean- sean- left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Contributor

@sean- sean- left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[review comment was posted twice for some reason]

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants