cmd/cli: fix log init end marker with partial data

For partial init log data (does not end with a newline), the log writer
discard data after the last newline to make the log prettier, then write
the init end marker. This causes the marker could be written more than
once, since the second overflows will preserve the data which does
include the marker from the first write.

To fix this, ensure that the init end marker is only written once, and
the second overflows will preserve data until the marker instead of the
fixed initial size like the first one.

cmd/cli/control_server.go

@@ -250,7 +250,7 @@ func (p *prog) registerControlServerHandler() {
 		}
 	}))
 	p.cs.register(sendLogsPath, http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
-		if time.Since(p.internalLogSent) < logSentInterval {
+		if time.Since(p.internalLogSent) < logWriterSentInterval {
 			w.WriteHeader(http.StatusServiceUnavailable)
 			return
 		}

cmd/cli/log_writer.go

@@ -16,13 +16,12 @@ import (
 )
 
 const (
-	logWriterSize        = 1024 * 1024 * 5 // 5 MB
-	logWriterSmallSize   = 1024 * 1024 * 1 // 1 MB
-	logWriterInitialSize = 32 * 1024       // 32 KB
-	logSentInterval      = time.Minute
-	logStartEndMarker    = "\n\n=== INIT_END ===\n\n"
-	logLogEndMarker      = "\n\n=== LOG_END ===\n\n"
-	logWarnEndMarker     = "\n\n=== WARN_END ===\n\n"
+	logWriterSize          = 1024 * 1024 * 5 // 5 MB
+	logWriterSmallSize     = 1024 * 1024 * 1 // 1 MB
+	logWriterInitialSize   = 32 * 1024       // 32 KB
+	logWriterSentInterval  = time.Minute
+	logWriterInitEndMarker = "\n\n=== INIT_END ===\n\n"
+	logWriterLogEndMarker  = "\n\n=== LOG_END ===\n\n"
 )
 
 type logViewResponse struct {
@@ -69,13 +68,23 @@ func (lw *logWriter) Write(p []byte) (int, error) {
 	// If writing p causes overflows, discard old data.
 	if lw.buf.Len()+len(p) > lw.size {
 		buf := lw.buf.Bytes()
-		buf = buf[:logWriterInitialSize]
-		if idx := bytes.LastIndex(buf, []byte("\n")); idx != -1 {
-			buf = buf[:idx]
+		haveEndMarker := false
+		// If there's init end marker already, preserve the data til the marker.
+		if idx := bytes.LastIndex(buf, []byte(logWriterInitEndMarker)); idx >= 0 {
+			buf = buf[:idx+len(logWriterInitEndMarker)]
+			haveEndMarker = true
+		} else {
+			// Otherwise, preserve the initial size data.
+			buf = buf[:logWriterInitialSize]
+			if idx := bytes.LastIndex(buf, []byte("\n")); idx != -1 {
+				buf = buf[:idx]
+			}
 		}
 		lw.buf.Reset()
 		lw.buf.Write(buf)
-		lw.buf.WriteString(logStartEndMarker) // indicate that the log was truncated.
+		if !haveEndMarker {
+			lw.buf.WriteString(logWriterInitEndMarker) // indicate that the log was truncated.
+		}
 	}
 	// If p is bigger than buffer size, truncate p by half until its size is smaller.
 	for len(p)+lw.buf.Len() > lw.size {
@@ -92,7 +101,7 @@ func (p *prog) initInternalLogging(writers []io.Writer) {
 	p.initInternalLogWriterOnce.Do(func() {
 		mainLog.Load().Notice().Msg("internal logging enabled")
 		p.internalLogWriter = newLogWriter()
-		p.internalLogSent = time.Now().Add(-logSentInterval)
+		p.internalLogSent = time.Now().Add(-logWriterSentInterval)
 		p.internalWarnLogWriter = newSmallLogWriter()
 	})
 	p.mu.Lock()
@@ -158,7 +167,7 @@ func (p *prog) logReader() (*logReader, error) {
 		wlwReader := bytes.NewReader(wlw.buf.Bytes())
 		wlwSize := wlw.buf.Len()
 		wlw.mu.Unlock()
-		reader := io.MultiReader(lwReader, bytes.NewReader([]byte(logLogEndMarker)), wlwReader)
+		reader := io.MultiReader(lwReader, bytes.NewReader([]byte(logWriterLogEndMarker)), wlwReader)
 		lr := &logReader{r: io.NopCloser(reader)}
 		lr.size = int64(lwSize + wlwSize)
 		if lr.size == 0 {

cmd/cli/log_writer_test.go

@@ -16,7 +16,7 @@ func Test_logWriter_Write(t *testing.T) {
 		t.Fatalf("unexpected buf content: %v", lw.buf.String())
 	}
 	newData := "B"
-	halfData := strings.Repeat("A", len(data)/2) + logStartEndMarker
+	halfData := strings.Repeat("A", len(data)/2) + logWriterInitEndMarker
 	lw.Write([]byte(newData))
 	if lw.buf.String() != halfData+newData {
 		t.Fatalf("unexpected new buf content: %v", lw.buf.String())
@@ -47,3 +47,39 @@ func Test_logWriter_ConcurrentWrite(t *testing.T) {
 		t.Fatalf("unexpected buf size: %v, content: %q", lw.buf.Len(), lw.buf.String())
 	}
 }
+
+func Test_logWriter_MarkerInitEnd(t *testing.T) {
+	size := 64 * 1024
+	lw := &logWriter{size: size}
+	lw.buf.Grow(lw.size)
+
+	paddingSize := 10
+	// Writing half of the size, minus len(end marker) and padding size.
+	dataSize := size/2 - len(logWriterInitEndMarker) - paddingSize
+	data := strings.Repeat("A", dataSize)
+	// Inserting newline for making partial init data
+	data += "\n"
+	// Filling left over buffer to make the log full.
+	// The data length: len(end marker) + padding size - 1 (for newline above) + size/2
+	data += strings.Repeat("A", len(logWriterInitEndMarker)+paddingSize-1+(size/2))
+	lw.Write([]byte(data))
+	if lw.buf.String() != data {
+		t.Fatalf("unexpected buf content: %v", lw.buf.String())
+	}
+	lw.Write([]byte("B"))
+	lw.Write([]byte(strings.Repeat("B", 256*1024)))
+	firstIdx := strings.Index(lw.buf.String(), logWriterInitEndMarker)
+	lastIdx := strings.LastIndex(lw.buf.String(), logWriterInitEndMarker)
+	// Check if init end marker present.
+	if firstIdx == -1 || lastIdx == -1 {
+		t.Fatalf("missing init end marker: %s", lw.buf.String())
+	}
+	// Check if init end marker appears only once.
+	if firstIdx != lastIdx {
+		t.Fatalf("log init end marker appears more than once: %s", lw.buf.String())
+	}
+	// Ensure that we have the correct init log data.
+	if !strings.Contains(lw.buf.String(), strings.Repeat("A", dataSize)+logWriterInitEndMarker) {
+		t.Fatalf("unexpected log content: %s", lw.buf.String())
+	}
+}