VictoriaMetrics/vendor/github.com/go-logfmt/logfmt/encode.go
Aliaksandr Valialkin 7f4fb34182 app/vmctl: move vmctl code from github.com/VictoriaMetrics/vmctl
It is better developing vmctl tool in VictoriaMetrics repository, so it could be released
together with the rest of vmutils tools such as vmalert, vmagent, vmbackup, vmrestore and vmauth.
2021-02-01 01:18:39 +02:00

322 lines
7.8 KiB
Go

package logfmt
import (
"bytes"
"encoding"
"errors"
"fmt"
"io"
"reflect"
"strings"
"unicode/utf8"
)
// MarshalKeyvals returns the logfmt encoding of keyvals, a variadic sequence
// of alternating keys and values.
func MarshalKeyvals(keyvals ...interface{}) ([]byte, error) {
buf := &bytes.Buffer{}
if err := NewEncoder(buf).EncodeKeyvals(keyvals...); err != nil {
return nil, err
}
return buf.Bytes(), nil
}
// An Encoder writes logfmt data to an output stream.
type Encoder struct {
w io.Writer
scratch bytes.Buffer
needSep bool
}
// NewEncoder returns a new encoder that writes to w.
func NewEncoder(w io.Writer) *Encoder {
return &Encoder{
w: w,
}
}
var (
space = []byte(" ")
equals = []byte("=")
newline = []byte("\n")
null = []byte("null")
)
// EncodeKeyval writes the logfmt encoding of key and value to the stream. A
// single space is written before the second and subsequent keys in a record.
// Nothing is written if a non-nil error is returned.
func (enc *Encoder) EncodeKeyval(key, value interface{}) error {
enc.scratch.Reset()
if enc.needSep {
if _, err := enc.scratch.Write(space); err != nil {
return err
}
}
if err := writeKey(&enc.scratch, key); err != nil {
return err
}
if _, err := enc.scratch.Write(equals); err != nil {
return err
}
if err := writeValue(&enc.scratch, value); err != nil {
return err
}
_, err := enc.w.Write(enc.scratch.Bytes())
enc.needSep = true
return err
}
// EncodeKeyvals writes the logfmt encoding of keyvals to the stream. Keyvals
// is a variadic sequence of alternating keys and values. Keys of unsupported
// type are skipped along with their corresponding value. Values of
// unsupported type or that cause a MarshalerError are replaced by their error
// but do not cause EncodeKeyvals to return an error. If a non-nil error is
// returned some key/value pairs may not have be written.
func (enc *Encoder) EncodeKeyvals(keyvals ...interface{}) error {
if len(keyvals) == 0 {
return nil
}
if len(keyvals)%2 == 1 {
keyvals = append(keyvals, nil)
}
for i := 0; i < len(keyvals); i += 2 {
k, v := keyvals[i], keyvals[i+1]
err := enc.EncodeKeyval(k, v)
if err == ErrUnsupportedKeyType {
continue
}
if _, ok := err.(*MarshalerError); ok || err == ErrUnsupportedValueType {
v = err
err = enc.EncodeKeyval(k, v)
}
if err != nil {
return err
}
}
return nil
}
// MarshalerError represents an error encountered while marshaling a value.
type MarshalerError struct {
Type reflect.Type
Err error
}
func (e *MarshalerError) Error() string {
return "error marshaling value of type " + e.Type.String() + ": " + e.Err.Error()
}
// ErrNilKey is returned by Marshal functions and Encoder methods if a key is
// a nil interface or pointer value.
var ErrNilKey = errors.New("nil key")
// ErrInvalidKey is returned by Marshal functions and Encoder methods if, after
// dropping invalid runes, a key is empty.
var ErrInvalidKey = errors.New("invalid key")
// ErrUnsupportedKeyType is returned by Encoder methods if a key has an
// unsupported type.
var ErrUnsupportedKeyType = errors.New("unsupported key type")
// ErrUnsupportedValueType is returned by Encoder methods if a value has an
// unsupported type.
var ErrUnsupportedValueType = errors.New("unsupported value type")
func writeKey(w io.Writer, key interface{}) error {
if key == nil {
return ErrNilKey
}
switch k := key.(type) {
case string:
return writeStringKey(w, k)
case []byte:
if k == nil {
return ErrNilKey
}
return writeBytesKey(w, k)
case encoding.TextMarshaler:
kb, err := safeMarshal(k)
if err != nil {
return err
}
if kb == nil {
return ErrNilKey
}
return writeBytesKey(w, kb)
case fmt.Stringer:
ks, ok := safeString(k)
if !ok {
return ErrNilKey
}
return writeStringKey(w, ks)
default:
rkey := reflect.ValueOf(key)
switch rkey.Kind() {
case reflect.Array, reflect.Chan, reflect.Func, reflect.Map, reflect.Slice, reflect.Struct:
return ErrUnsupportedKeyType
case reflect.Ptr:
if rkey.IsNil() {
return ErrNilKey
}
return writeKey(w, rkey.Elem().Interface())
}
return writeStringKey(w, fmt.Sprint(k))
}
}
// keyRuneFilter returns r for all valid key runes, and -1 for all invalid key
// runes. When used as the mapping function for strings.Map and bytes.Map
// functions it causes them to remove invalid key runes from strings or byte
// slices respectively.
func keyRuneFilter(r rune) rune {
if r <= ' ' || r == '=' || r == '"' || r == utf8.RuneError {
return -1
}
return r
}
func writeStringKey(w io.Writer, key string) error {
k := strings.Map(keyRuneFilter, key)
if k == "" {
return ErrInvalidKey
}
_, err := io.WriteString(w, k)
return err
}
func writeBytesKey(w io.Writer, key []byte) error {
k := bytes.Map(keyRuneFilter, key)
if len(k) == 0 {
return ErrInvalidKey
}
_, err := w.Write(k)
return err
}
func writeValue(w io.Writer, value interface{}) error {
switch v := value.(type) {
case nil:
return writeBytesValue(w, null)
case string:
return writeStringValue(w, v, true)
case []byte:
return writeBytesValue(w, v)
case encoding.TextMarshaler:
vb, err := safeMarshal(v)
if err != nil {
return err
}
if vb == nil {
vb = null
}
return writeBytesValue(w, vb)
case error:
se, ok := safeError(v)
return writeStringValue(w, se, ok)
case fmt.Stringer:
ss, ok := safeString(v)
return writeStringValue(w, ss, ok)
default:
rvalue := reflect.ValueOf(value)
switch rvalue.Kind() {
case reflect.Array, reflect.Chan, reflect.Func, reflect.Map, reflect.Slice, reflect.Struct:
return ErrUnsupportedValueType
case reflect.Ptr:
if rvalue.IsNil() {
return writeBytesValue(w, null)
}
return writeValue(w, rvalue.Elem().Interface())
}
return writeStringValue(w, fmt.Sprint(v), true)
}
}
func needsQuotedValueRune(r rune) bool {
return r <= ' ' || r == '=' || r == '"' || r == utf8.RuneError
}
func writeStringValue(w io.Writer, value string, ok bool) error {
var err error
if ok && value == "null" {
_, err = io.WriteString(w, `"null"`)
} else if strings.IndexFunc(value, needsQuotedValueRune) != -1 {
_, err = writeQuotedString(w, value)
} else {
_, err = io.WriteString(w, value)
}
return err
}
func writeBytesValue(w io.Writer, value []byte) error {
var err error
if bytes.IndexFunc(value, needsQuotedValueRune) != -1 {
_, err = writeQuotedBytes(w, value)
} else {
_, err = w.Write(value)
}
return err
}
// EndRecord writes a newline character to the stream and resets the encoder
// to the beginning of a new record.
func (enc *Encoder) EndRecord() error {
_, err := enc.w.Write(newline)
if err == nil {
enc.needSep = false
}
return err
}
// Reset resets the encoder to the beginning of a new record.
func (enc *Encoder) Reset() {
enc.needSep = false
}
func safeError(err error) (s string, ok bool) {
defer func() {
if panicVal := recover(); panicVal != nil {
if v := reflect.ValueOf(err); v.Kind() == reflect.Ptr && v.IsNil() {
s, ok = "null", false
} else {
s, ok = fmt.Sprintf("PANIC:%v", panicVal), false
}
}
}()
s, ok = err.Error(), true
return
}
func safeString(str fmt.Stringer) (s string, ok bool) {
defer func() {
if panicVal := recover(); panicVal != nil {
if v := reflect.ValueOf(str); v.Kind() == reflect.Ptr && v.IsNil() {
s, ok = "null", false
} else {
s, ok = fmt.Sprintf("PANIC:%v", panicVal), true
}
}
}()
s, ok = str.String(), true
return
}
func safeMarshal(tm encoding.TextMarshaler) (b []byte, err error) {
defer func() {
if panicVal := recover(); panicVal != nil {
if v := reflect.ValueOf(tm); v.Kind() == reflect.Ptr && v.IsNil() {
b, err = nil, nil
} else {
b, err = nil, fmt.Errorf("panic when marshalling: %s", panicVal)
}
}
}()
b, err = tm.MarshalText()
if err != nil {
return nil, &MarshalerError{
Type: reflect.TypeOf(tm),
Err: err,
}
}
return
}