feat: preserve order when writing json

charmbracelet · May 7, 2024 · dd7bc94 · dd7bc94
1 parent deae1b0
commit dd7bc94
Show file tree

Hide file tree

Showing 2 changed files with 213 additions and 16 deletions.
diff --git a/json.go b/json.go
@@ -1,34 +1,36 @@
 package log
 
 import (
+	"bytes"
 	"encoding/json"
 	"fmt"
 	"time"
 )
 
 func (l *Logger) jsonFormatter(keyvals ...interface{}) {
-	m := make(map[string]interface{}, len(keyvals)/2)
+	jw := jsonWriter{w: &l.b}
+	jw.start()
 	for i := 0; i < len(keyvals); i += 2 {
 		switch keyvals[i] {
 		case TimestampKey:
 			if t, ok := keyvals[i+1].(time.Time); ok {
-				m[TimestampKey] = t.Format(l.timeFormat)
+				jw.write(TimestampKey, t.Format(l.timeFormat))
 			}
 		case LevelKey:
 			if level, ok := keyvals[i+1].(Level); ok {
-				m[LevelKey] = level.String()
+				jw.write(LevelKey, level.String())
 			}
 		case CallerKey:
 			if caller, ok := keyvals[i+1].(string); ok {
-				m[CallerKey] = caller
+				jw.write(CallerKey, caller)
 			}
 		case PrefixKey:
 			if prefix, ok := keyvals[i+1].(string); ok {
-				m[PrefixKey] = prefix
+				jw.write(PrefixKey, prefix)
 			}
 		case MessageKey:
 			if msg := keyvals[i+1]; msg != nil {
-				m[MessageKey] = fmt.Sprint(msg)
+				jw.write(MessageKey, fmt.Sprint(msg))
 			}
 		default:
 			var (
@@ -51,11 +53,60 @@ func (l *Logger) jsonFormatter(keyvals ...interface{}) {
 			default:
 				val = v
 			}
-			m[key] = val
+			jw.write(key, val)
 		}
 	}
+	jw.end()
+}
+
+type jsonWriter struct {
+	w *bytes.Buffer
+}
+
+func (w *jsonWriter) start() {
+	w.w.WriteRune('{')
+}
+
+func (w *jsonWriter) end() {
+	w.w.WriteRune('}')
+	w.w.WriteRune('\n')
+}
+
+func (w *jsonWriter) write(key string, value any) {
+	// store pos if we need to rewind
+	pos := w.w.Len()
+
+	// add separator when buffer is longer than '{'
+	if w.w.Len() > 1 {
+		w.w.WriteRune(',')
+	}
+
+	err := w.writeEncoded(key)
+	if err != nil {
+		w.w.Truncate(pos)
+		return
+	}
+	w.w.WriteRune(':')
 
-	e := json.NewEncoder(&l.b)
+	pos = w.w.Len()
+	err = w.writeEncoded(value)
+	if err != nil {
+		w.w.Truncate(pos)
+		w.w.WriteString(`"invalid value"`)
+	}
+}
+
+func (w *jsonWriter) writeEncoded(v any) error {
+	e := json.NewEncoder(w.w)
 	e.SetEscapeHTML(false)
-	_ = e.Encode(m)
+	if err := e.Encode(v); err != nil {
+		return err
+	}
+
+	// trailing \n added by json.Encode
+	b := w.w.Bytes()
+	if len(b) > 0 && b[len(b)-1] == '\n' {
+		w.w.Truncate(w.w.Len() - 1)
+	}
+	return nil
 }
diff --git a/json_test.go b/json_test.go
@@ -59,14 +59,14 @@ func TestJson(t *testing.T) {
 		},
 		{
 			name:     "odd number of kvs",
-			expected: "{\"baz\":\"missing value\",\"foo\":\"bar\",\"level\":\"error\",\"msg\":\"info\"}\n",
+			expected: "{\"level\":\"error\",\"msg\":\"info\",\"foo\":\"bar\",\"baz\":\"missing value\"}\n",
 			msg:      "info",
 			kvs:      []interface{}{"foo", "bar", "baz"},
 			f:        l.Error,
 		},
 		{
 			name:     "error field",
-			expected: "{\"error\":\"error message\",\"level\":\"error\",\"msg\":\"info\"}\n",
+			expected: "{\"level\":\"error\",\"msg\":\"info\",\"error\":\"error message\"}\n",
 			msg:      "info",
 			kvs:      []interface{}{"error", errors.New("error message")},
 			f:        l.Error,
@@ -108,7 +108,7 @@ func TestJson(t *testing.T) {
 		},
 		{
 			name:     "map of strings",
-			expected: "{\"level\":\"info\",\"map\":{\"a\":\"b\",\"foo\":\"bar\"},\"msg\":\"info\"}\n",
+			expected: "{\"level\":\"info\",\"msg\":\"info\",\"map\":{\"a\":\"b\",\"foo\":\"bar\"}}\n",
 			msg:      "info",
 			kvs:      []interface{}{"map", map[string]string{"a": "b", "foo": "bar"}},
 			f:        l.Info,
@@ -140,14 +140,14 @@ func TestJsonCaller(t *testing.T) {
 	}{
 		{
 			name:     "simple caller",
-			expected: fmt.Sprintf("{\"caller\":\"log/%s:%d\",\"level\":\"info\",\"msg\":\"info\"}\n", filepath.Base(file), line+30),
+			expected: fmt.Sprintf("{\"level\":\"info\",\"caller\":\"log/%s:%d\",\"msg\":\"info\"}\n", filepath.Base(file), line+30),
 			msg:      "info",
 			kvs:      nil,
 			f:        l.Info,
 		},
 		{
 			name:     "nested caller",
-			expected: fmt.Sprintf("{\"caller\":\"log/%s:%d\",\"level\":\"info\",\"msg\":\"info\"}\n", filepath.Base(file), line+30),
+			expected: fmt.Sprintf("{\"level\":\"info\",\"caller\":\"log/%s:%d\",\"msg\":\"info\"}\n", filepath.Base(file), line+30),
 			msg:      "info",
 			kvs:      nil,
 			f: func(msg interface{}, kvs ...interface{}) {
@@ -165,17 +165,163 @@ func TestJsonCaller(t *testing.T) {
 	}
 }
 
+func TestJsonTime(t *testing.T) {
+	var buf bytes.Buffer
+	logger := New(&buf)
+	logger.SetTimeFunction(_zeroTime)
+	logger.SetFormatter(JSONFormatter)
+	logger.SetReportTimestamp(true)
+	logger.Info("info")
+	require.Equal(t, "{\"time\":\"0002/01/01 00:00:00\",\"level\":\"info\",\"msg\":\"info\"}\n", buf.String())
+}
+
+func TestJsonPrefix(t *testing.T) {
+	var buf bytes.Buffer
+	logger := New(&buf)
+	logger.SetFormatter(JSONFormatter)
+	logger.SetPrefix("my-prefix")
+	logger.Info("info")
+	require.Equal(t, "{\"level\":\"info\",\"prefix\":\"my-prefix\",\"msg\":\"info\"}\n", buf.String())
+}
+
 func TestJsonCustomKey(t *testing.T) {
 	var buf bytes.Buffer
 	oldTsKey := TimestampKey
 	defer func() {
 		TimestampKey = oldTsKey
 	}()
-	TimestampKey = "time"
+	TimestampKey = "other-time"
 	logger := New(&buf)
 	logger.SetTimeFunction(_zeroTime)
 	logger.SetFormatter(JSONFormatter)
 	logger.SetReportTimestamp(true)
 	logger.Info("info")
-	require.Equal(t, "{\"level\":\"info\",\"msg\":\"info\",\"time\":\"0002/01/01 00:00:00\"}\n", buf.String())
+	require.Equal(t, "{\"other-time\":\"0002/01/01 00:00:00\",\"level\":\"info\",\"msg\":\"info\"}\n", buf.String())
+}
+
+func TestJsonWriter(t *testing.T) {
+	testCases := []struct {
+		name     string
+		fn       func(w *jsonWriter)
+		expected string
+	}{
+		{
+			"string",
+			func(w *jsonWriter) {
+				w.start()
+				w.write("a", "value")
+				w.end()
+			},
+			`{"a":"value"}` + "\n",
+		},
+		{
+			"int",
+			func(w *jsonWriter) {
+				w.start()
+				w.write("a", 123)
+				w.end()
+			},
+			`{"a":123}` + "\n",
+		},
+		{
+			"bytes",
+			func(w *jsonWriter) {
+				w.start()
+				w.write("b", []byte{0x0, 0x1})
+				w.end()
+			},
+			`{"b":"AAE="}` + "\n",
+		},
+		{
+			"no fields",
+			func(w *jsonWriter) {
+				w.start()
+				w.end()
+			},
+			`{}` + "\n",
+		},
+		{
+			"multiple in asc order",
+			func(w *jsonWriter) {
+				w.start()
+				w.write("a", "value")
+				w.write("b", "some-other")
+				w.end()
+			},
+			`{"a":"value","b":"some-other"}` + "\n",
+		},
+		{
+			"multiple in desc order",
+			func(w *jsonWriter) {
+				w.start()
+				w.write("b", "some-other")
+				w.write("a", "value")
+				w.end()
+			},
+			`{"b":"some-other","a":"value"}` + "\n",
+		},
+		{
+			"depth",
+			func(w *jsonWriter) {
+				w.start()
+				w.write("a", map[string]int{"b": 123})
+				w.end()
+			},
+			`{"a":{"b":123}}` + "\n",
+		},
+		{
+			"key contains reserved",
+			func(w *jsonWriter) {
+				w.start()
+				w.write("a:\"b", "value")
+				w.end()
+			},
+			`{"a:\"b":"value"}` + "\n",
+		},
+		{
+			"pointer",
+			func(w *jsonWriter) {
+				w.start()
+				w.write("a", ptr("pointer"))
+				w.end()
+			},
+			`{"a":"pointer"}` + "\n",
+		},
+		{
+			"double-pointer",
+			func(w *jsonWriter) {
+				w.start()
+				w.write("a", ptr(ptr("pointer")))
+				w.end()
+			},
+			`{"a":"pointer"}` + "\n",
+		},
+		{
+			"invalid",
+			func(w *jsonWriter) {
+				w.start()
+				w.write("a", invalidJSON{})
+				w.end()
+			},
+			`{"a":"invalid value"}` + "\n",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			var buf bytes.Buffer
+			tc.fn(&jsonWriter{w: &buf})
+			require.Equal(t, tc.expected, buf.String())
+		})
+	}
+}
+
+func ptr[T any](v T) *T {
+	return &v
+}
+
+type invalidJSON struct{}
+
+func (invalidJSON) MarshalJSON() ([]byte, error) {
+	return nil, errors.New("invalid json error")
 }