Config setting to only yield complete lines (#26)

* first draft of CompleteLines

* tests

* refactor

* CompleteLines should return the last line if we don't follow

* rename temp dirs to avoid races

* Conditional allocation of strings.Builder

* Fix tests by using the new cleanup()

Co-authored-by: nxadm <[email protected]>

Author: kokes

tail.go

@@ -77,8 +77,9 @@ type Config struct {
 	Pipe      bool      // The file is a named pipe (mkfifo)
 
 	// Generic IO
-	Follow      bool // Continue looking for new lines (tail -f)
-	MaxLineSize int  // If non-zero, split longer lines into multiple lines
+	Follow        bool // Continue looking for new lines (tail -f)
+	MaxLineSize   int  // If non-zero, split longer lines into multiple lines
+	CompleteLines bool // Only return complete lines (that end with "\n" or EOF when Follow is false)
 
 	// Optionally, use a ratelimiter (e.g. created by the ratelimiter/NewLeakyBucket function)
 	RateLimiter *ratelimiter.LeakyBucket
@@ -97,6 +98,8 @@ type Tail struct {
 	reader  *bufio.Reader
 	lineNum int
 
+	lineBuf *strings.Builder
+
 	watcher watch.FileWatcher
 	changes *watch.FileChanges
 
@@ -128,6 +131,10 @@ func TailFile(filename string, config Config) (*Tail, error) {
 		Config:   config,
 	}
 
+	if config.CompleteLines {
+		t.lineBuf = new(strings.Builder)
+	}
+
 	// when Logger was not specified in config, use default logger
 	if t.Logger == nil {
 		t.Logger = DefaultLogger
@@ -202,6 +209,9 @@ func (tail *Tail) closeFile() {
 }
 
 func (tail *Tail) reopen() error {
+	if tail.lineBuf != nil {
+		tail.lineBuf.Reset()
+	}
 	tail.closeFile()
 	tail.lineNum = 0
 	for {
@@ -229,16 +239,32 @@ func (tail *Tail) readLine() (string, error) {
 	tail.lk.Lock()
 	line, err := tail.reader.ReadString('\n')
 	tail.lk.Unlock()
-	if err != nil {
+
+	newlineEnding := strings.HasSuffix(line, "\n")
+	line = strings.TrimRight(line, "\n")
+
+	// if we don't have to handle incomplete lines, we can return the line as-is
+	if !tail.Config.CompleteLines {
 		// Note ReadString "returns the data read before the error" in
 		// case of an error, including EOF, so we return it as is. The
 		// caller is expected to process it if err is EOF.
 		return line, err
 	}
 
-	line = strings.TrimRight(line, "\n")
+	if _, err := tail.lineBuf.WriteString(line); err != nil {
+		return line, err
+	}
 
-	return line, err
+	if newlineEnding {
+		line = tail.lineBuf.String()
+		tail.lineBuf.Reset()
+		return line, nil
+	} else {
+		if tail.Config.Follow {
+			line = ""
+		}
+		return line, io.EOF
+	}
 }
 
 func (tail *Tail) tailFileSync() {

tail_test.go

@@ -515,6 +515,113 @@ func TestInotify_WaitForCreateThenMove(t *testing.T) {
 	tailTest.Cleanup(tail, false)
 }
 
+func TestIncompleteLines(t *testing.T) {
+	tailTest, cleanup := NewTailTest("incomplete-lines", t)
+	defer cleanup()
+	filename := "test.txt"
+	config := Config{
+		Follow:        true,
+		CompleteLines: true,
+	}
+	tail := tailTest.StartTail(filename, config)
+	go func() {
+		time.Sleep(100 * time.Millisecond)
+		tailTest.CreateFile(filename, "hello world\n")
+		time.Sleep(100 * time.Millisecond)
+		// here we intentially write a partial line to see if `Tail` contains
+		// information that it's incomplete
+		tailTest.AppendFile(filename, "hello")
+		time.Sleep(100 * time.Millisecond)
+		tailTest.AppendFile(filename, " again\n")
+	}()
+
+	lines := []string{"hello world", "hello again"}
+
+	tailTest.ReadLines(tail, lines, false)
+
+	tailTest.RemoveFile(filename)
+	tail.Stop()
+	tail.Cleanup()
+}
+
+func TestIncompleteLongLines(t *testing.T) {
+	tailTest, cleanup := NewTailTest("incomplete-lines-long", t)
+	defer cleanup()
+	filename := "test.txt"
+	config := Config{
+		Follow:        true,
+		MaxLineSize:   3,
+		CompleteLines: true,
+	}
+	tail := tailTest.StartTail(filename, config)
+	go func() {
+		time.Sleep(100 * time.Millisecond)
+		tailTest.CreateFile(filename, "hello world\n")
+		time.Sleep(100 * time.Millisecond)
+		tailTest.AppendFile(filename, "hello")
+		time.Sleep(100 * time.Millisecond)
+		tailTest.AppendFile(filename, "again\n")
+	}()
+
+	lines := []string{"hel", "lo ", "wor", "ld", "hel", "loa", "gai", "n"}
+
+	tailTest.ReadLines(tail, lines, false)
+
+	tailTest.RemoveFile(filename)
+	tail.Stop()
+	tail.Cleanup()
+}
+
+func TestIncompleteLinesWithReopens(t *testing.T) {
+	tailTest, cleanup := NewTailTest("incomplete-lines-reopens", t)
+	defer cleanup()
+	filename := "test.txt"
+	config := Config{
+		Follow:        true,
+		CompleteLines: true,
+	}
+	tail := tailTest.StartTail(filename, config)
+	go func() {
+		time.Sleep(100 * time.Millisecond)
+		tailTest.CreateFile(filename, "hello world\nhi")
+		time.Sleep(100 * time.Millisecond)
+		tailTest.TruncateFile(filename, "rewriting\n")
+	}()
+
+	// not that the "hi" gets lost, because it was never a complete line
+	lines := []string{"hello world", "rewriting"}
+
+	tailTest.ReadLines(tail, lines, false)
+
+	tailTest.RemoveFile(filename)
+	tail.Stop()
+	tail.Cleanup()
+}
+
+func TestIncompleteLinesWithoutFollow(t *testing.T) {
+	tailTest, cleanup := NewTailTest("incomplete-lines-no-follow", t)
+	defer cleanup()
+	filename := "test.txt"
+	config := Config{
+		Follow:        false,
+		CompleteLines: true,
+	}
+	tail := tailTest.StartTail(filename, config)
+	go func() {
+		time.Sleep(100 * time.Millisecond)
+		// intentionally missing a newline at the end
+		tailTest.CreateFile(filename, "foo\nbar\nbaz")
+	}()
+
+	lines := []string{"foo", "bar", "baz"}
+
+	tailTest.VerifyTailOutput(tail, lines, true)
+
+	tailTest.RemoveFile(filename)
+	tail.Stop()
+	tail.Cleanup()
+}
+
 func reSeek(t *testing.T, poll bool) {
 	var name string
 	if poll {
@@ -557,7 +664,7 @@ type TailTest struct {
 }
 
 func NewTailTest(name string, t *testing.T) (TailTest, func()) {
-	testdir, err := ioutil.TempDir("", "tail-test-" + name)
+	testdir, err := ioutil.TempDir("", "tail-test-"+name)
 	if err != nil {
 		t.Fatal(err)
 	}