diff --git a/common/hugio/hasBytesWriter.go b/common/hugio/hasBytesWriter.go index d2bcd1bb4..c22376ca7 100644 --- a/common/hugio/hasBytesWriter.go +++ b/common/hugio/hasBytesWriter.go @@ -21,8 +21,9 @@ import ( type HasBytesWriter struct { Patterns []*HasBytesPattern - i int done bool + // The tail of the bytes written so far, retained so we can detect a + // pattern that straddles the boundary between two Write calls. buff []byte } @@ -31,10 +32,13 @@ type HasBytesPattern struct { Pattern []byte } -func (h *HasBytesWriter) patternLen() int { +// maxPatternLen returns the length of the longest pattern. +func (h *HasBytesWriter) maxPatternLen() int { l := 0 for _, p := range h.Patterns { - l += len(p.Pattern) + if len(p.Pattern) > l { + l = len(p.Pattern) + } } return l } @@ -44,36 +48,55 @@ func (h *HasBytesWriter) Write(p []byte) (n int, err error) { return len(p), nil } - if len(h.buff) == 0 { - h.buff = make([]byte, h.patternLen()*2) + keep := h.maxPatternLen() - 1 + + // Join the tail retained from previous Writes with the head of this chunk + // so a pattern straddling the boundary is still detected. Only the + // boundary window is copied; the chunk itself is scanned in place below. + var boundary []byte + if keep > 0 && len(h.buff) > 0 { + head := p + if len(head) > keep { + head = head[:keep] + } + boundary = make([]byte, 0, len(h.buff)+len(head)) + boundary = append(boundary, h.buff...) + boundary = append(boundary, head...) } - for i := range p { - h.buff[h.i] = p[i] - h.i++ - if h.i == len(h.buff) { - // Shift left. - copy(h.buff, h.buff[len(h.buff)/2:]) - h.i = len(h.buff) / 2 + // Scan each not-yet-matched pattern once per Write instead of once per byte. + done := true + for _, pp := range h.Patterns { + if pp.Match { + continue } - - for _, pp := range h.Patterns { - if bytes.Contains(h.buff, pp.Pattern) { - pp.Match = true - done := true - for _, ppp := range h.Patterns { - if !ppp.Match { - done = false - break - } - } - if done { - h.done = true - } - return len(p), nil - } + if bytes.Contains(p, pp.Pattern) || bytes.Contains(boundary, pp.Pattern) { + pp.Match = true + continue } + done = false + } + if done { + // All patterns found; no need to look at any more data. + h.done = true + h.buff = nil + return len(p), nil + } + + // Retain the last keep bytes of (previous tail + this chunk) to detect a + // pattern straddling into the next Write. + switch { + case keep <= 0: + h.buff = h.buff[:0] + case len(p) >= keep: + h.buff = append(h.buff[:0], p[len(p)-keep:]...) + default: + // Chunk shorter than keep: slide the window over the retained tail. + if total := len(h.buff) + len(p); total > keep { + h.buff = h.buff[total-keep:] + } + h.buff = append(h.buff, p...) } return len(p), nil diff --git a/common/hugio/hasBytesWriter_test.go b/common/hugio/hasBytesWriter_test.go index 9e689a112..009d29b29 100644 --- a/common/hugio/hasBytesWriter_test.go +++ b/common/hugio/hasBytesWriter_test.go @@ -65,3 +65,61 @@ func TestHasBytesWriter(t *testing.T) { fmt.Fprintf(w, "__foo") c.Assert(h.Patterns[0].Match, qt.Equals, true) } + +func TestHasBytesWriterMultiplePatterns(t *testing.T) { + c := qt.New(t) + + neww := func() (*HasBytesWriter, io.Writer) { + var b bytes.Buffer + h := &HasBytesWriter{ + Patterns: []*HasBytesPattern{ + {Pattern: []byte("__hdeferred/")}, + {Pattern: []byte("__h_pp_l1")}, + }, + } + return h, io.MultiWriter(&b, h) + } + + // Neither pattern present. + h, w := neww() + fmt.Fprint(w, "the quick brown fox jumps over the lazy dog") + c.Assert(h.Patterns[0].Match, qt.Equals, false) + c.Assert(h.Patterns[1].Match, qt.Equals, false) + c.Assert(h.done, qt.Equals, false) + + // Only the second pattern present; the writer must not report a match + // for the first, and must not prematurely mark itself done. + h, w = neww() + fmt.Fprint(w, "prefix __h_pp_l1 suffix") + c.Assert(h.Patterns[0].Match, qt.Equals, false) + c.Assert(h.Patterns[1].Match, qt.Equals, true) + c.Assert(h.done, qt.Equals, false) + + // Both patterns present across multiple writes; done once all match. + h, w = neww() + fmt.Fprint(w, "aaa __hdef") + fmt.Fprint(w, "erred/xyz bbb __h_p") + fmt.Fprint(w, "p_l1 ccc") + c.Assert(h.Patterns[0].Match, qt.Equals, true) + c.Assert(h.Patterns[1].Match, qt.Equals, true) + c.Assert(h.done, qt.Equals, true) +} + +func BenchmarkHasBytesWriter(b *testing.B) { + // A large chunk of output containing neither pattern is the common case + // (a normal rendered page): the writer must scan all of it. + content := []byte(strings.Repeat("
baz
\n", 4000)) + + b.ResetTimer() + for range b.N { + h := &HasBytesWriter{ + Patterns: []*HasBytesPattern{ + {Pattern: []byte("__hdeferred/")}, + {Pattern: []byte("__h_pp_l1")}, + }, + } + if _, err := h.Write(content); err != nil { + b.Fatal(err) + } + } +}