Improve SQL query building (#137)

* metastore: Extract SQL query building and add benchmark

* metastore: Optimize query building for retrieving location lines

* metastore: Use improved SQL query pre-allocation heuristic

Instead of pre-allocating pessimistically for the largest possible
integer, the new heuristic uses the largest integer that actually occurs
in the query.

pkg/storage/metastore/sql.go

@@ -19,6 +19,7 @@ import (
 	"fmt"
 	"strconv"
 	"time"
+	"unsafe"
 
 	"github.com/google/pprof/profile"
 	"go.opentelemetry.io/otel/attribute"
@@ -386,6 +387,7 @@ func (s *sqlMetaStore) getLinesByLocationIDs(ctx context.Context, ids ...uint64)
 
 	res := make(map[uint64][]locationLine, len(ids))
 	remainingIds := []uint64{}
+	maxRemainingId := uint64(0)
 	for _, id := range ids {
 		ll, found, err := s.cache.getLocationLinesByID(ctx, id)
 		if err != nil {
@@ -401,6 +403,9 @@ func (s *sqlMetaStore) getLinesByLocationIDs(ctx context.Context, ids ...uint64)
 			res[id] = ll
 			continue
 		}
+		if maxRemainingId < id {
+			maxRemainingId = id
+		}
 		remainingIds = append(remainingIds, id)
 	}
 	ids = remainingIds
@@ -409,19 +414,7 @@ func (s *sqlMetaStore) getLinesByLocationIDs(ctx context.Context, ids ...uint64)
 		return res, functionIDs, nil
 	}
 
-	sIds := ""
-	for i, id := range ids {
-		if i > 0 {
-			sIds += ","
-		}
-		sIds += strconv.FormatInt(int64(id), 10)
-	}
-
-	rows, err := s.db.QueryContext(ctx,
-		fmt.Sprintf(
-			`SELECT "location_id", "line", "function_id"
-				FROM "lines" WHERE location_id IN (%s)`, sIds),
-	)
+	rows, err := s.db.QueryContext(ctx, buildLinesByLocationIDsQuery(maxRemainingId, ids))
 	if err != nil {
 		return nil, nil, fmt.Errorf("execute SQL query: %w", err)
 	}
@@ -470,6 +463,41 @@ func (s *sqlMetaStore) getLinesByLocationIDs(ctx context.Context, ids ...uint64)
 	return res, functionIDs, nil
 }
 
+const (
+	queryStart = `SELECT "location_id", "line", "function_id"
+			FROM "lines" WHERE location_id IN (`
+	comma          = ','
+	closingBracket = ')'
+)
+
+func buildLinesByLocationIDsQuery(max uint64, ids []uint64) string {
+	maxLen := len(strconv.FormatUint(max, 10))
+
+	query := make([]byte, 0,
+		// The max value is known, and invididual string can be larger than it.
+		len(ids)*maxLen+
+			// Add the start of the query.
+			len(queryStart)+
+			// len(ids)-1 commas, and a closing bracket is len(ids).
+			len(ids),
+	)
+	query = append(query, queryStart...)
+
+	for i := range ids {
+		if i > 0 {
+			query = append(query, comma)
+		}
+		query = strconv.AppendUint(query, ids[i], 10)
+	}
+
+	query = append(query, closingBracket)
+	return unsafeString(query)
+}
+
+func unsafeString(b []byte) string {
+	return *((*string)(unsafe.Pointer(&b)))
+}
+
 func (s *sqlMetaStore) getFunctionsByIDs(ctx context.Context, ids ...uint64) (map[uint64]*profile.Function, error) {
 	ctx, span := s.tracer.Start(ctx, "getFunctionsByIDs")
 	defer span.End()

pkg/storage/metastore/sql_bench_test.go

@@ -0,0 +1,38 @@
+// Copyright 2021 The Parca Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package metastore
+
+import (
+	"fmt"
+	"math"
+	"testing"
+)
+
+var result string
+
+func BenchmarkBuildLinesByLocationIDsQuery(b *testing.B) {
+	for k := 0.; k <= 6; k++ {
+		n := uint64(math.Pow(10, k))
+		b.Run(fmt.Sprintf("%d", n), func(b *testing.B) {
+			input := make([]uint64, 0, n)
+			for i := uint64(0); i < n; i++ {
+				input = append(input, i)
+			}
+
+			for i := 0; i < b.N; i++ {
+				result = buildLinesByLocationIDsQuery(input[len(input)-1], input)
+			}
+		})
+	}
+}