riverlog: build capped log arrays in one pass

Appending to large river:log histories still did extra work in the sjson
version: it built an appended array first, then repeatedly deleted index 0
while over cap. In high-volume histories this caused multiple full-array
rewrites before reaching the final payload.

Rework append/prune to compute the kept suffix first and then build the
final array once. The new flow marshals only the new entry, derives bounds
for existing array elements, decides how many oldest entries to drop to fit
MaxTotalBytes, and emits a single final JSON array. It also truncates log
bytes before converting to string so oversized log buffers are not copied
in full.

Benchmark comparison for `BenchmarkMiddlewareWorkAppend`
(`-benchtime=100ms`, same machine):

case: 256 KB
impl            ns/op      B/op  allocs/op
before         821178   1167325         25
after          177355    884454         31

case: 2 MB
impl            ns/op      B/op  allocs/op
before        4854185   6479606         29
after         1247496   7006818         46

Added focused unit coverage for appendLogDataWithCap, including missing-key
behavior, non-array metadata errors, minimal oldest-first pruning, keeping
newest entry when over cap, and no-cap behavior.

Author: bgentry

riverlog/river_log.go

@@ -13,7 +13,6 @@ import (
 	"log/slog"
 
 	"github.com/tidwall/gjson"
-	"github.com/tidwall/sjson"
 
 	"github.com/riverqueue/river/internal/jobexecutor"
 	"github.com/riverqueue/river/rivershared/baseservice"
@@ -186,26 +185,26 @@ func (m *Middleware) Work(ctx context.Context, job *rivertype.JobRow, doInner fu
 
 	// This all runs invariant of whether the job panics or returns an error.
 	defer func() {
-		logData := logBuf.String()
+		logBytes := logBuf.Bytes()
 
 		// Return early if nothing ended up getting logged.
-		if len(logData) < 1 {
+		if len(logBytes) < 1 {
 			return
 		}
 
 		// Postgres JSONB is limited to 255MB, but it would be a bad idea to get
 		// anywhere close to that limit here.
-		if len(logData) > m.config.MaxSizeBytes {
+		if len(logBytes) > m.config.MaxSizeBytes {
 			m.Logger.WarnContext(ctx, m.Name+": Logs size exceeded maximum; truncating",
-				slog.Int("logs_size", len(logData)),
+				slog.Int("logs_size", len(logBytes)),
 				slog.Int("max_size", m.config.MaxSizeBytes),
 			)
-			logData = logData[0:m.config.MaxSizeBytes]
+			logBytes = logBytes[0:m.config.MaxSizeBytes]
 		}
 
 		newLogEntryBytes, err := json.Marshal(logAttempt{
 			Attempt: job.Attempt,
-			Log:     logData,
+			Log:     string(logBytes),
 		})
 		if err != nil {
 			m.Logger.ErrorContext(ctx, m.Name+": Error marshaling log data",
@@ -236,42 +235,125 @@ func (m *Middleware) Work(ctx context.Context, job *rivertype.JobRow, doInner fu
 }
 
 func appendLogDataWithCap(metadataBytes, newLogEntryBytes []byte, maxTotalBytes int) ([]byte, int, error) {
-	if !json.Valid(metadataBytes) {
-		return nil, 0, errors.New("metadata is not valid JSON")
-	}
-
 	existingLogData := gjson.GetBytes(metadataBytes, metadataKey)
-	var allLogDataBytes []byte
+	var existingLogArrayBytes []byte
 	switch {
 	case !existingLogData.Exists():
-		allLogDataBytes = []byte("[]")
+		existingLogArrayBytes = []byte("[]")
 	case existingLogData.IsArray():
-		allLogDataBytes = []byte(existingLogData.Raw)
+		// Slice raw JSON straight from metadata bytes to avoid an extra copy.
+		existingLogArrayBytes = metadataBytes[existingLogData.Index : existingLogData.Index+len(existingLogData.Raw)]
 	default:
 		return nil, 0, fmt.Errorf("%q value is not an array", metadataKey)
 	}
 
-	var err error
-	allLogDataBytes, err = sjson.SetRawBytes(allLogDataBytes, "-1", newLogEntryBytes)
+	existingElementBounds, err := getArrayElementBounds(existingLogArrayBytes)
 	if err != nil {
 		return nil, 0, err
 	}
 
-	if maxTotalBytes <= 0 || len(allLogDataBytes) <= maxTotalBytes {
-		return allLogDataBytes, 0, nil
+	// Determine the smallest suffix to keep that still fits with the new entry.
+	// This keeps pruning oldest-first while avoiding repeated full rewrites.
+	keepStart := getKeepStart(existingElementBounds, len(newLogEntryBytes), maxTotalBytes)
+
+	// Build the final array once from the kept suffix plus the new entry.
+	appendedLogDataBytes := buildAppendedArray(existingLogArrayBytes, existingElementBounds, keepStart, newLogEntryBytes)
+	numDroppedEntries := keepStart
+	if numDroppedEntries > len(existingElementBounds) {
+		numDroppedEntries = len(existingElementBounds)
 	}
 
-	// Drop oldest entries first, while always retaining the latest one.
-	numEntries := len(gjson.ParseBytes(allLogDataBytes).Array())
-	var numDroppedEntries int
-	for numEntries > 1 && len(allLogDataBytes) > maxTotalBytes {
-		allLogDataBytes, err = sjson.DeleteBytes(allLogDataBytes, "0")
-		if err != nil {
-			return nil, numDroppedEntries, err
+	return appendedLogDataBytes, numDroppedEntries, nil
+}
+
+type arrayElementBounds struct {
+	Start int
+	End   int
+}
+
+func getArrayElementBounds(arrayBytes []byte) ([]arrayElementBounds, error) {
+	arrResult := gjson.ParseBytes(arrayBytes)
+	if !arrResult.IsArray() {
+		return nil, errors.New("expected a JSON array")
+	}
+
+	elements := arrResult.Array()
+	bounds := make([]arrayElementBounds, len(elements))
+	for i, elem := range elements {
+		if elem.Index < 0 {
+			return nil, errors.New("failed to determine array element index")
+		}
+		bounds[i] = arrayElementBounds{
+			Start: elem.Index,
+			End:   elem.Index + len(elem.Raw),
 		}
-		numEntries--
-		numDroppedEntries++
 	}
+	return bounds, nil
+}
+
+func getKeepStart(bounds []arrayElementBounds, newEntryLen, maxTotalBytes int) int {
+	if maxTotalBytes <= 0 {
+		return 0
+	}
+
+	// Keep newest entry even if it's larger than the configured cap.
+	newOnlyLen := 2 + newEntryLen // `[` + entry + `]`
+	if newOnlyLen > maxTotalBytes {
+		return len(bounds)
+	}
+
+	if len(bounds) == 0 {
+		return 0
+	}
+
+	// suffixContentLen[i] is the serialized content length for elements i..end
+	// inside the array, including commas between elements.
+	suffixContentLen := make([]int, len(bounds))
+	for i := len(bounds) - 1; i >= 0; i-- {
+		elemLen := bounds[i].End - bounds[i].Start
+		if i == len(bounds)-1 {
+			suffixContentLen[i] = elemLen
+		} else {
+			suffixContentLen[i] = elemLen + 1 + suffixContentLen[i+1]
+		}
+	}
+
+	// Iterate from oldest to newest so we drop the minimum number of entries
+	// necessary to fit the configured cap.
+	for keepStart := 0; keepStart <= len(bounds); keepStart++ {
+		contentLen := newEntryLen
+		if keepStart < len(bounds) {
+			contentLen += 1 + suffixContentLen[keepStart] // comma between kept suffix and new entry
+		}
+		totalLen := 2 + contentLen // `[` + content + `]`
+		if totalLen <= maxTotalBytes {
+			return keepStart
+		}
+	}
+
+	return len(bounds)
+}
+
+func buildAppendedArray(existingArrayBytes []byte, bounds []arrayElementBounds, keepStart int, newEntryBytes []byte) []byte {
+	var totalLen int
+	if keepStart >= len(bounds) {
+		totalLen = 2 + len(newEntryBytes)
+	} else {
+		// Kept suffix is contiguous in the original array bytes, so copy once.
+		keptContentLen := bounds[len(bounds)-1].End - bounds[keepStart].Start
+		totalLen = 2 + keptContentLen + 1 + len(newEntryBytes) // [] + comma + new entry
+	}
+
+	result := make([]byte, 0, totalLen)
+	result = append(result, '[')
+
+	if keepStart < len(bounds) {
+		result = append(result, existingArrayBytes[bounds[keepStart].Start:bounds[len(bounds)-1].End]...)
+		result = append(result, ',')
+	}
+
+	result = append(result, newEntryBytes...)
+	result = append(result, ']')
 
-	return allLogDataBytes, numDroppedEntries, nil
+	return result
 }

riverlog/river_log_test.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"log"
 	"log/slog"
+	"strings"
 	"testing"
 
 	"github.com/jackc/pgx/v5"
@@ -44,6 +45,112 @@ func TestLogger(t *testing.T) {
 	})
 }
 
+func TestAppendLogDataWithCap(t *testing.T) {
+	t.Parallel()
+
+	marshalLog := func(tb testing.TB, attempt int, log string) []byte {
+		tb.Helper()
+		b, err := json.Marshal(logAttempt{Attempt: attempt, Log: log})
+		require.NoError(tb, err)
+		return b
+	}
+
+	marshalMetadataWithLogs := func(tb testing.TB, logs []logAttempt) []byte {
+		tb.Helper()
+		b, err := json.Marshal(map[string]any{
+			metadataKey: logs,
+		})
+		require.NoError(tb, err)
+		return b
+	}
+
+	unmarshalLogs := func(tb testing.TB, rawArray []byte) []logAttempt {
+		tb.Helper()
+		var logs []logAttempt
+		require.NoError(tb, json.Unmarshal(rawArray, &logs))
+		return logs
+	}
+
+	t.Run("MissingKeyStartsFromEmptyArray", func(t *testing.T) {
+		t.Parallel()
+
+		newEntry := marshalLog(t, 1, "new")
+		result, dropped, err := appendLogDataWithCap([]byte(`{"other":"value"}`), newEntry, maxTotalBytes)
+		require.NoError(t, err)
+		require.Zero(t, dropped)
+		require.Equal(t, []logAttempt{{Attempt: 1, Log: "new"}}, unmarshalLogs(t, result))
+	})
+
+	t.Run("NonArrayLogValueReturnsError", func(t *testing.T) {
+		t.Parallel()
+
+		newEntry := marshalLog(t, 1, "new")
+		_, _, err := appendLogDataWithCap([]byte(`{"river:log":{"not":"array"}}`), newEntry, maxTotalBytes)
+		require.EqualError(t, err, `"river:log" value is not an array`)
+	})
+
+	t.Run("PrunesOldestEntriesOnlyAsNeeded", func(t *testing.T) {
+		t.Parallel()
+
+		existing := []logAttempt{
+			{Attempt: 1, Log: "a"},
+			{Attempt: 2, Log: "b"},
+			{Attempt: 3, Log: "c"},
+		}
+		newEntry := marshalLog(t, 4, "d")
+
+		target, err := json.Marshal([]logAttempt{
+			{Attempt: 2, Log: "b"},
+			{Attempt: 3, Log: "c"},
+			{Attempt: 4, Log: "d"},
+		})
+		require.NoError(t, err)
+
+		result, dropped, err := appendLogDataWithCap(marshalMetadataWithLogs(t, existing), newEntry, len(target))
+		require.NoError(t, err)
+		require.Equal(t, 1, dropped)
+		require.Equal(t, []logAttempt{
+			{Attempt: 2, Log: "b"},
+			{Attempt: 3, Log: "c"},
+			{Attempt: 4, Log: "d"},
+		}, unmarshalLogs(t, result))
+	})
+
+	t.Run("KeepsNewestEntryEvenIfOverCap", func(t *testing.T) {
+		t.Parallel()
+
+		existing := []logAttempt{
+			{Attempt: 1, Log: "a"},
+			{Attempt: 2, Log: "b"},
+		}
+		newEntry := marshalLog(t, 3, strings.Repeat("x", 64))
+
+		result, dropped, err := appendLogDataWithCap(marshalMetadataWithLogs(t, existing), newEntry, 8)
+		require.NoError(t, err)
+		require.Equal(t, len(existing), dropped)
+		require.Equal(t, []logAttempt{
+			{Attempt: 3, Log: strings.Repeat("x", 64)},
+		}, unmarshalLogs(t, result))
+	})
+
+	t.Run("NoCapKeepsEverything", func(t *testing.T) {
+		t.Parallel()
+
+		existing := []logAttempt{
+			{Attempt: 1, Log: "a"},
+		}
+		newEntry := marshalLog(t, 2, "b")
+
+		result, dropped, err := appendLogDataWithCap(marshalMetadataWithLogs(t, existing), newEntry, 0)
+		require.NoError(t, err)
+		require.Zero(t, dropped)
+		require.Equal(t, []logAttempt{
+			{Attempt: 1, Log: "a"},
+			{Attempt: 2, Log: "b"},
+		}, unmarshalLogs(t, result))
+	})
+}
+
 func TestLoggerSafely(t *testing.T) {
 	t.Parallel()