Backport querylog fixes

go/logstats/logger.go

@@ -0,0 +1,217 @@
+/*
+Copyright 2024 The Vitess 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 logstats
+
+import (
+	"io"
+	"slices"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
+	"vitess.io/vitess/go/hack"
+	"vitess.io/vitess/go/sqltypes"
+	querypb "vitess.io/vitess/go/vt/proto/query"
+)
+
+type logbv struct {
+	Name string
+	BVar *querypb.BindVariable
+}
+
+// Logger is a zero-allocation logger for logstats.
+// It can output logs as JSON or as plaintext, following the commonly used
+// logstats format that is shared between the tablets and the gates.
+type Logger struct {
+	b     []byte
+	bvars []logbv
+	n     int
+	json  bool
+}
+
+func sortBVars(sorted []logbv, bvars map[string]*querypb.BindVariable) []logbv {
+	for k, bv := range bvars {
+		sorted = append(sorted, logbv{k, bv})
+	}
+	slices.SortFunc(sorted, func(a, b logbv) int {
+		return strings.Compare(a.Name, b.Name)
+	})
+	return sorted
+}
+
+func (log *Logger) appendBVarsJSON(b []byte, bvars map[string]*querypb.BindVariable, full bool) []byte {
+	log.bvars = sortBVars(log.bvars[:0], bvars)
+
+	b = append(b, '{')
+	for i, bv := range log.bvars {
+		if i > 0 {
+			b = append(b, ',', ' ')
+		}
+		b = strconv.AppendQuote(b, bv.Name)
+		b = append(b, `: {"type": `...)
+		b = strconv.AppendQuote(b, querypb.Type_name[int32(bv.BVar.Type)])
+		b = append(b, `, "value": `...)
+
+		if sqltypes.IsIntegral(bv.BVar.Type) || sqltypes.IsFloat(bv.BVar.Type) {
+			b = append(b, bv.BVar.Value...)
+		} else if bv.BVar.Type == sqltypes.Tuple {
+			b = append(b, '"')
+			b = strconv.AppendInt(b, int64(len(bv.BVar.Values)), 10)
+			b = append(b, ` items"`...)
+		} else {
+			if full {
+				b = strconv.AppendQuote(b, hack.String(bv.BVar.Value))
+			} else {
+				b = append(b, '"')
+				b = strconv.AppendInt(b, int64(len(bv.BVar.Values)), 10)
+				b = append(b, ` bytes"`...)
+			}
+		}
+		b = append(b, '}')
+	}
+	return append(b, '}')
+}
+
+func (log *Logger) Init(json bool) {
+	log.n = 0
+	log.json = json
+	if log.json {
+		log.b = append(log.b, '{')
+	}
+}
+
+func (log *Logger) Redacted() {
+	log.String("[REDACTED]")
+}
+
+func (log *Logger) Key(key string) {
+	if log.json {
+		if log.n > 0 {
+			log.b = append(log.b, ',', ' ')
+		}
+		log.b = append(log.b, '"')
+		log.b = append(log.b, key...)
+		log.b = append(log.b, '"', ':', ' ')
+	} else {
+		if log.n > 0 {
+			log.b = append(log.b, '\t')
+		}
+	}
+	log.n++
+}
+
+func (log *Logger) StringUnquoted(value string) {
+	if log.json {
+		log.b = strconv.AppendQuote(log.b, value)
+	} else {
+		log.b = append(log.b, value...)
+	}
+}
+
+func (log *Logger) TabTerminated() {
+	if !log.json {
+		log.b = append(log.b, '\t')
+	}
+}
+
+func (log *Logger) String(value string) {
+	log.b = strconv.AppendQuote(log.b, value)
+}
+
+func (log *Logger) StringSingleQuoted(value string) {
+	if log.json {
+		log.b = strconv.AppendQuote(log.b, value)
+	} else {
+		log.b = append(log.b, '\'')
+		log.b = append(log.b, value...)
+		log.b = append(log.b, '\'')
+	}
+}
+
+func (log *Logger) Time(t time.Time) {
+	const timeFormat = "2006-01-02 15:04:05.000000"
+	if log.json {
+		log.b = append(log.b, '"')
+		log.b = t.AppendFormat(log.b, timeFormat)
+		log.b = append(log.b, '"')
+	} else {
+		log.b = t.AppendFormat(log.b, timeFormat)
+	}
+}
+
+func (log *Logger) Duration(t time.Duration) {
+	log.b = strconv.AppendFloat(log.b, t.Seconds(), 'f', 6, 64)
+}
+
+func (log *Logger) BindVariables(bvars map[string]*querypb.BindVariable, full bool) {
+	// the bind variables are printed as JSON in text mode because the original
+	// printing syntax, which was simply `fmt.Sprintf("%v")`, is not stable or
+	// safe to parse
+	log.b = log.appendBVarsJSON(log.b, bvars, full)
+}
+
+func (log *Logger) Int(i int64) {
+	log.b = strconv.AppendInt(log.b, i, 10)
+}
+
+func (log *Logger) Uint(u uint64) {
+	log.b = strconv.AppendUint(log.b, u, 10)
+}
+
+func (log *Logger) Bool(b bool) {
+	log.b = strconv.AppendBool(log.b, b)
+}
+
+func (log *Logger) Strings(strs []string) {
+	log.b = append(log.b, '[')
+	for i, t := range strs {
+		if i > 0 {
+			log.b = append(log.b, ',')
+		}
+		log.b = strconv.AppendQuote(log.b, t)
+	}
+	log.b = append(log.b, ']')
+}
+
+func (log *Logger) Flush(w io.Writer) (err error) {
+	if log.json {
+		log.b = append(log.b, '}')
+	}
+	log.b = append(log.b, '\n')
+	_, err = w.Write(log.b)
+
+	clear(log.bvars)
+	log.bvars = log.bvars[:0]
+	log.b = log.b[:0]
+	log.n = 0
+
+	loggerPool.Put(log)
+	return err
+}
+
+var loggerPool = sync.Pool{New: func() any {
+	return &Logger{}
+}}
+
+// NewLogger returns a new Logger instance to perform logstats logging.
+// The logger must be initialized with (*Logger).Init before usage and
+// flushed with (*Logger).Flush once all the key-values have been written
+// to it.
+func NewLogger() *Logger {
+	return loggerPool.Get().(*Logger)
+}

go/vt/vtgate/logstats/logstats.go

@@ -18,21 +18,16 @@ package logstats
 
 import (
 	"context"
-	"encoding/json"
-	"fmt"
 	"io"
 	"net/url"
 	"time"
 
 	"github.com/google/safehtml"
 
-	"vitess.io/vitess/go/sqltypes"
+	"vitess.io/vitess/go/logstats"
 	"vitess.io/vitess/go/streamlog"
-	"vitess.io/vitess/go/tb"
 	"vitess.io/vitess/go/vt/callerid"
 	"vitess.io/vitess/go/vt/callinfo"
-	"vitess.io/vitess/go/vt/log"
-
 	querypb "vitess.io/vitess/go/vt/proto/query"
 )
 
@@ -128,69 +123,60 @@ func (stats *LogStats) Logf(w io.Writer, params url.Values) error {
 		return nil
 	}
 
-	// FormatBindVariables call might panic so we're going to catch it here
-	// and print out the stack trace for debugging.
-	defer func() {
-		if x := recover(); x != nil {
-			log.Errorf("Uncaught panic:\n%v\n%s", x, tb.Stack(4))
-		}
-	}()
-
-	formattedBindVars := "\"[REDACTED]\""
-	if !streamlog.GetRedactDebugUIQueries() {
-		_, fullBindParams := params["full"]
-		formattedBindVars = sqltypes.FormatBindVariables(
-			stats.BindVariables,
-			fullBindParams,
-			streamlog.GetQueryLogFormat() == streamlog.QueryLogFormatJSON,
-		)
-	}
-
-	// TODO: remove username here we fully enforce immediate caller id
+	redacted := streamlog.GetRedactDebugUIQueries()
+	_, fullBindParams := params["full"]
 	remoteAddr, username := stats.RemoteAddrUsername()
 
-	var fmtString string
-	switch streamlog.GetQueryLogFormat() {
-	case streamlog.QueryLogFormatText:
-		fmtString = "%v\t%v\t%v\t'%v'\t'%v'\t%v\t%v\t%.6f\t%.6f\t%.6f\t%.6f\t%v\t%q\t%v\t%v\t%v\t%q\t%q\t%q\t%v\t%v\t%q\n"
-	case streamlog.QueryLogFormatJSON:
-		fmtString = "{\"Method\": %q, \"RemoteAddr\": %q, \"Username\": %q, \"ImmediateCaller\": %q, \"Effective Caller\": %q, \"Start\": \"%v\", \"End\": \"%v\", \"TotalTime\": %.6f, \"PlanTime\": %v, \"ExecuteTime\": %v, \"CommitTime\": %v, \"StmtType\": %q, \"SQL\": %q, \"BindVars\": %v, \"ShardQueries\": %v, \"RowsAffected\": %v, \"Error\": %q, \"TabletType\": %q, \"SessionUUID\": %q, \"Cached Plan\": %v, \"TablesUsed\": %v, \"ActiveKeyspace\": %q}\n"
-	}
-
-	tables := stats.TablesUsed
-	if tables == nil {
-		tables = []string{}
-	}
-	tablesUsed, marshalErr := json.Marshal(tables)
-	if marshalErr != nil {
-		return marshalErr
+	log := logstats.NewLogger()
+	log.Init(streamlog.GetQueryLogFormat() == streamlog.QueryLogFormatJSON)
+	log.Key("Method")
+	log.StringUnquoted(stats.Method)
+	log.Key("RemoteAddr")
+	log.StringUnquoted(remoteAddr)
+	log.Key("Username")
+	log.StringUnquoted(username)
+	log.Key("ImmediateCaller")
+	log.StringSingleQuoted(stats.ImmediateCaller())
+	log.Key("Effective Caller")
+	log.StringSingleQuoted(stats.EffectiveCaller())
+	log.Key("Start")
+	log.Time(stats.StartTime)
+	log.Key("End")
+	log.Time(stats.EndTime)
+	log.Key("TotalTime")
+	log.Duration(stats.TotalTime())
+	log.Key("PlanTime")
+	log.Duration(stats.PlanTime)
+	log.Key("ExecuteTime")
+	log.Duration(stats.ExecuteTime)
+	log.Key("CommitTime")
+	log.Duration(stats.CommitTime)
+	log.Key("StmtType")
+	log.StringUnquoted(stats.StmtType)
+	log.Key("SQL")
+	log.String(stats.SQL)
+	log.Key("BindVars")
+	if redacted {
+		log.Redacted()
+	} else {
+		log.BindVariables(stats.BindVariables, fullBindParams)
 	}
-	_, err := fmt.Fprintf(
-		w,
-		fmtString,
-		stats.Method,
-		remoteAddr,
-		username,
-		stats.ImmediateCaller(),
-		stats.EffectiveCaller(),
-		stats.StartTime.Format("2006-01-02 15:04:05.000000"),
-		stats.EndTime.Format("2006-01-02 15:04:05.000000"),
-		stats.TotalTime().Seconds(),
-		stats.PlanTime.Seconds(),
-		stats.ExecuteTime.Seconds(),
-		stats.CommitTime.Seconds(),
-		stats.StmtType,
-		stats.SQL,
-		formattedBindVars,
-		stats.ShardQueries,
-		stats.RowsAffected,
-		stats.ErrorStr(),
-		stats.TabletType,
-		stats.SessionUUID,
-		stats.CachedPlan,
-		string(tablesUsed),
-		stats.ActiveKeyspace,
-	)
-
-	return err
+	log.Key("ShardQueries")
+	log.Uint(stats.ShardQueries)
+	log.Key("RowsAffected")
+	log.Uint(stats.RowsAffected)
+	log.Key("Error")
+	log.String(stats.ErrorStr())
+	log.Key("TabletType")
+	log.String(stats.TabletType)
+	log.Key("SessionUUID")
+	log.String(stats.SessionUUID)
+	log.Key("Cached Plan")
+	log.Bool(stats.CachedPlan)
+	log.Key("TablesUsed")
+	log.Strings(stats.TablesUsed)
+	log.Key("ActiveKeyspace")
+	log.String(stats.ActiveKeyspace)
+
+	return log.Flush(w)
 }