+package main
+
+import (
+ "bufio"
+ "errors"
+ "fmt"
+ "os"
+ "os/exec"
+ "strconv"
+ "strings"
+ "sync"
+ "testing"
+ "time"
+
+ "github.com/stretchr/testify/require"
+)
+
+// childDirVariable names the chat directory of the writer child process of
+// TestKilledWriterLosesNoLoggedMessage.
+const childDirVariable = "ONECHAT_TEST_CHILD_DIR"
+
+func TestFlushWaitsForASyncThatStartsAfterTheWrite(t *testing.T) {
+ gate := newGatedSync()
+ s := newSyncer(gate.Sync)
+ s.Written()
+ gate.WaitStarted(t, 1)
+
+ // This write comes during the first sync, so that sync does not cover
+ // it.
+ s.Written()
+ flushed := make(chan error)
+ go func() {
+ flushed <- s.Flush(t.Context())
+ }()
+
+ gate.Release()
+ gate.WaitStarted(t, 2)
+ select {
+ case err := <-flushed:
+ t.Fatalf("flush returned before the second sync: %v", err)
+ case <-time.After(20 * time.Millisecond):
+ }
+
+ gate.Release()
+ require.NoError(t, <-flushed)
+ require.NoError(t, s.Close())
+ require.Equal(t, 2, gate.Count())
+}
+
+func TestOneSyncCoversTheWritesDuringTheSyncBeforeIt(t *testing.T) {
+ gate := newGatedSync()
+ s := newSyncer(gate.Sync)
+ s.Written()
+ gate.WaitStarted(t, 1)
+ for range 10 {
+ s.Written()
+ }
+
+ gate.Release()
+ gate.WaitStarted(t, 2)
+ gate.Release()
+ require.NoError(t, s.Flush(t.Context()))
+ require.Equal(t, 2, gate.Count())
+ require.NoError(t, s.Close())
+}
+
+func TestFlushWithoutWritesNeedsNoSync(t *testing.T) {
+ gate := newGatedSync()
+ s := newSyncer(gate.Sync)
+ require.NoError(t, s.Flush(t.Context()))
+ require.NoError(t, s.Close())
+ require.Equal(t, 0, gate.Count())
+}
+
+func TestAFailedSyncFailsEachLaterFlush(t *testing.T) {
+ s := newSyncer(func() error {
+ return errors.New("disk error")
+ })
+
+ s.Written()
+ require.ErrorContains(t, s.Flush(t.Context()), "disk error")
+
+ s.Written()
+ require.ErrorContains(t, s.Flush(t.Context()), "disk error")
+ require.ErrorContains(t, s.Close(), "disk error")
+}
+
+func TestMemoryFlushSyncsTheChat(t *testing.T) {
+ memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer())
+ require.NoError(t, err)
+
+ _, err = memory.Append("tool", "shell: ls")
+ require.NoError(t, err)
+ require.NoError(t, memory.Flush(t.Context()))
+ require.NoError(t, memory.Close())
+}
+
+func TestKilledWriterLosesNoLoggedMessage(t *testing.T) {
+ root := t.TempDir()
+ for round := range 3 {
+ acknowledged := runKilledWriter(t, root, 200+100*round)
+
+ memory, err := openMemory(t.Context(), root, time.Now, newFakeSummarizer())
+ require.NoError(t, err)
+ require.GreaterOrEqual(t, memory.log.Len(), acknowledged)
+ for id := range memory.log.Len() {
+ message, err := memory.log.Message(id)
+ require.NoError(t, err)
+ require.Equal(t, writerKind(id), message.Kind)
+ require.Equal(t, writerText(id), message.Text)
+ }
+
+ require.NoError(t, memory.Close())
+ }
+}
+
+// TestWriterChild is the child process of TestKilledWriterLosesNoLoggedMessage.
+// It appends messages to the chat in the directory of childDirVariable and
+// prints the ID of each message after its Append returns, until it is killed.
+func TestWriterChild(t *testing.T) {
+ root := os.Getenv(childDirVariable)
+ if root == "" {
+ t.Skip("runs only as a child process")
+ }
+
+ memory, err := openMemory(t.Context(), root, time.Now, newFakeSummarizer())
+ require.NoError(t, err)
+
+ for id := memory.log.Len(); ; id++ {
+ _, err = memory.Append(writerKind(id), writerText(id))
+ require.NoError(t, err)
+
+ _, err = fmt.Println(id)
+ require.NoError(t, err)
+ }
+}
+
+// gatedSync is a sync function that blocks until the test releases it.
+type gatedSync struct {
+ mutex sync.Mutex
+ count int
+ started chan struct{}
+ release chan struct{}
+}
+
+func newGatedSync() *gatedSync {
+ return &gatedSync{
+ started: make(chan struct{}, 100),
+ release: make(chan struct{}),
+ }
+}
+
+func (g *gatedSync) Sync() error {
+ g.mutex.Lock()
+ g.count++
+ g.mutex.Unlock()
+
+ g.started <- struct{}{}
+ <-g.release
+
+ return nil
+}
+
+// WaitStarted waits until sync number n has started.
+func (g *gatedSync) WaitStarted(t *testing.T, n int) {
+ t.Helper()
+
+ require.Eventually(t, func() bool {
+ return g.Count() >= n
+ }, time.Second, time.Millisecond)
+}
+
+// Release lets the running sync end.
+func (g *gatedSync) Release() {
+ <-g.started
+ g.release <- struct{}{}
+}
+
+func (g *gatedSync) Count() int {
+ g.mutex.Lock()
+ defer g.mutex.Unlock()
+
+ return g.count
+}
+
+// runKilledWriter starts the writer child on the chat in root, kills it with
+// SIGKILL after it acknowledged count messages, and returns the number of
+// messages that it acknowledged.
+func runKilledWriter(t *testing.T, root string, count int) int {
+ t.Helper()
+
+ command := exec.Command(os.Args[0], "-test.run=^TestWriterChild$")
+ command.Env = append(os.Environ(), childDirVariable+"="+root)
+ stdout, err := command.StdoutPipe()
+ require.NoError(t, err)
+ require.NoError(t, command.Start())
+
+ acknowledged := 0
+ scanner := bufio.NewScanner(stdout)
+ for acknowledged < count && scanner.Scan() {
+ _, err = strconv.Atoi(strings.TrimSpace(scanner.Text()))
+ if err == nil {
+ acknowledged++
+ }
+ }
+
+ require.NoError(t, command.Process.Kill())
+ for scanner.Scan() {
+ _, err = strconv.Atoi(strings.TrimSpace(scanner.Text()))
+ if err == nil {
+ acknowledged++
+ }
+ }
+
+ require.Error(t, command.Wait())
+ require.Equal(t, count, min(acknowledged, count), "the child stopped early")
+
+ return acknowledged
+}
+
+// writerKind returns the kind of message id of the writer child.
+func writerKind(id int) string {
+ return []string{"user", "tool", "echo", "reply"}[id%4]
+}
+
+// writerText returns the text of message id of the writer child. Some texts
+// are large, so a kill can come during the write of one line.
+func writerText(id int) string {
+ size := 40 + id%7*30
+ if id%5 == 4 {
+ size = 64*1024 + id
+ }
+
+ return strings.Repeat(strconv.Itoa(id%10), size)
+}
projects / onechat / commitdiff