Buffer N bytes from the host and show to clients

Clients can see empty screen when they join the hosts. This buffers the
last N bytes from the host and show to the clients when they first join.

go.mod

@@ -50,6 +50,7 @@ require (
 require (
 	github.com/eiannone/keyboard v0.0.0-20220611211555-0d226195f203
 	github.com/google/go-github/v48 v48.2.0
+	github.com/stretchr/testify v1.7.0
 	golang.org/x/exp v0.0.0-20220407100705-7b9b53b0aca4
 	golang.org/x/term v0.12.0
 )
@@ -59,6 +60,7 @@ require (
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/cespare/xxhash/v2 v2.2.0 // indirect
 	github.com/cpuguy83/go-md2man/v2 v2.0.2 // indirect
+	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/fsnotify/fsnotify v1.6.0 // indirect
 	github.com/go-kit/log v0.2.1 // indirect
 	github.com/go-logfmt/logfmt v0.5.1 // indirect
@@ -73,6 +75,7 @@ require (
 	github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect
 	github.com/nu7hatch/gouuid v0.0.0-20131221200532-179d4d0c4d8d // indirect
 	github.com/pkg/errors v0.9.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/prometheus/client_model v0.4.1-0.20230718164431-9a2bf3000d16 // indirect
 	github.com/prometheus/common v0.44.0 // indirect
 	github.com/prometheus/procfs v0.11.1 // indirect

host/internal/server.go

@@ -37,7 +37,7 @@ type Server struct {
 }
 
 func (s *Server) ServeWithContext(ctx context.Context, l net.Listener) error {
-	writers := uio.NewMultiWriter()
+	writers := uio.NewMultiWriter(5)
 
 	cmdCtx, cmdCancel := context.WithCancel(ctx)
 	defer cmdCancel()

io/writer.go

@@ -5,40 +5,81 @@ import (
 	"sync"
 )
 
-func NewMultiWriter(writers ...io.Writer) *MultiWriter {
-	return &MultiWriter{writers: writers}
+type buffer struct {
+	mu sync.Mutex
+
+	data [][]byte
+	size int
+}
+
+func (c *buffer) Append(p []byte) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if len(c.data) >= c.size {
+		c.data = c.data[1:]
+	}
+
+	pp := make([]byte, len(p))
+	copy(pp, p)
+
+	c.data = append(c.data, pp)
+}
+
+func (c *buffer) Size() int {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	return len(c.data)
+}
+
+func (c *buffer) Data() [][]byte {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	result := make([][]byte, len(c.data))
+	return append(result, c.data...)
+}
+
+func NewMultiWriter(bufferSize int, writers ...io.Writer) *MultiWriter {
+	return &MultiWriter{
+		writers: writers,
+		buffer:  &buffer{size: bufferSize},
+	}
 }
 
 // MultiWriter is a concurrent safe writer that allows appending/removing writers.
 // Newly appended writers get the last write to preserve last output.
 type MultiWriter struct {
-	mu      sync.Mutex
+	writeMu sync.Mutex
 	writers []io.Writer
-	cache   []byte
+
+	buffer *buffer
 }
 
 func (t *MultiWriter) Append(writers ...io.Writer) error {
-	t.mu.Lock()
-	defer t.mu.Unlock()
-
-	// write last cache to new writers
-	if len(t.cache) > 0 {
+	// write last buffer to new writers
+	if t.buffer.Size() > 0 {
 		for _, w := range writers {
-			_, err := w.Write(t.cache)
-			if err != nil {
-				return err
+			for _, d := range t.buffer.Data() {
+				_, err := w.Write(d)
+				if err != nil {
+					return err
+				}
 			}
 		}
 	}
 
+	t.writeMu.Lock()
+	defer t.writeMu.Unlock()
 	t.writers = append(t.writers, writers...)
 
 	return nil
 }
 
 func (t *MultiWriter) Remove(writers ...io.Writer) {
-	t.mu.Lock()
-	defer t.mu.Unlock()
+	t.writeMu.Lock()
+	defer t.writeMu.Unlock()
 
 	for i := len(t.writers) - 1; i > 0; i-- {
 		for _, v := range writers {
@@ -51,12 +92,10 @@ func (t *MultiWriter) Remove(writers ...io.Writer) {
 }
 
 func (t *MultiWriter) Write(p []byte) (n int, err error) {
-	t.mu.Lock()
-	defer t.mu.Unlock()
+	t.buffer.Append(p)
 
-	// reset cache
-	t.cache = make([]byte, len(p))
-	copy(t.cache, p)
+	t.writeMu.Lock()
+	defer t.writeMu.Unlock()
 
 	for _, w := range t.writers {
 		n, err = w.Write(p)

io/writer_test.go

@@ -5,54 +5,45 @@ import (
 	"io"
 	"testing"
 
-	"github.com/google/go-cmp/cmp"
+	"github.com/stretchr/testify/assert"
 )
 
 func Test_MultiWriter(t *testing.T) {
+	assert := assert.New(t)
+
 	w1 := bytes.NewBuffer(nil)
-	w := NewMultiWriter(w1)
+	w := NewMultiWriter(1, w1)
 
 	r := bytes.NewBufferString("hello1")
 	_, _ = io.Copy(w, r)
 
-	want := "hello1"
-	got := w1.String()
-	if diff := cmp.Diff(want, got); diff != "" {
-		t.Errorf("want=%s got=%s:\n%s", want, got, diff)
-	}
+	assert.Equal("hello1", w1.String())
 
 	// append w2
 	r = bytes.NewBufferString("hello2")
 	w2 := bytes.NewBuffer(nil)
 	_ = w.Append(w2)
 	_, _ = io.Copy(w, r)
 
-	want = "hello1hello2"
-	got = w1.String()
-	if diff := cmp.Diff(want, got); diff != "" {
-		t.Errorf("want=%s got=%s:\n%s", want, got, diff)
-	}
+	assert.Equal("hello1hello2", w1.String())
+	assert.Equal("hello1hello2", w2.String())
 
-	want = "hello1hello2" // new writer has the last write
-	got = w2.String()
-	if diff := cmp.Diff(want, got); diff != "" {
-		t.Errorf("want=%s got=%s:\n%s", want, got, diff)
-	}
+	// append w3
+	r = bytes.NewBufferString("hello3")
+	w3 := bytes.NewBuffer(nil)
+	_ = w.Append(w3)
+	_, _ = io.Copy(w, r)
+
+	assert.Equal("hello1hello2hello3", w1.String())
+	assert.Equal("hello1hello2hello3", w2.String())
+	assert.Equal("hello2hello3", w3.String())
 
 	// remove w2
-	r = bytes.NewBufferString("hello3")
+	r = bytes.NewBufferString("hello4")
 	w.Remove(w2)
 	_, _ = io.Copy(w, r)
 
-	want = "hello1hello2hello3"
-	got = w1.String()
-	if diff := cmp.Diff(want, got); diff != "" {
-		t.Errorf("want=%s got=%s:\n%s", want, got, diff)
-	}
-
-	want = "hello1hello2" // removed writer doesn't have the latest write
-	got = w2.String()
-	if diff := cmp.Diff(want, got); diff != "" {
-		t.Errorf("want=%s got=%s:\n%s", want, got, diff)
-	}
+	assert.Equal("hello1hello2hello3hello4", w1.String())
+	assert.Equal("hello1hello2hello3", w2.String())
+	assert.Equal("hello2hello3hello4", w3.String())
 }