Merge branch 'main' into 15814-wb2-secrets
[arvados.git] / services / keep-web / writebuffer_test.go
1 // Copyright (C) The Arvados Authors. All rights reserved.
2 //
3 // SPDX-License-Identifier: AGPL-3.0
4
5 package keepweb
6
7 import (
8         "bytes"
9         "io"
10         "math/rand"
11         "time"
12
13         . "gopkg.in/check.v1"
14 )
15
16 var _ = Suite(&writeBufferSuite{})
17
18 type writeBufferSuite struct {
19 }
20
21 // 1000 / 96.3 ns/op = 10.384 GB/s
22 func (s *writeBufferSuite) Benchmark_1KBWrites(c *C) {
23         wb := newWriteBuffer(io.Discard, 1<<20)
24         in := make([]byte, 1000)
25         for i := 0; i < c.N; i++ {
26                 wb.Write(in)
27         }
28         wb.Close()
29 }
30
31 func (s *writeBufferSuite) TestRandomizedSpeedsAndSizes(c *C) {
32         for i := 0; i < 20; i++ {
33                 insize := rand.Intn(1 << 26)
34                 bufsize := rand.Intn(1 << 26)
35                 if i < 2 {
36                         // make sure to test edge cases
37                         bufsize = i
38                 } else if insize/bufsize > 1000 {
39                         // don't waste too much time testing tiny
40                         // buffer / huge content
41                         insize = bufsize*1000 + 123
42                 }
43                 c.Logf("%s: insize %d bufsize %d", c.TestName(), insize, bufsize)
44
45                 in := make([]byte, insize)
46                 b := byte(0)
47                 for i := range in {
48                         in[i] = b
49                         b++
50                 }
51
52                 out := &bytes.Buffer{}
53                 done := make(chan struct{})
54                 pr, pw := io.Pipe()
55                 go func() {
56                         n, err := slowCopy(out, pr, rand.Intn(8192)+1)
57                         c.Check(err, IsNil)
58                         c.Check(n, Equals, int64(insize))
59                         close(done)
60                 }()
61                 wb := newWriteBuffer(pw, bufsize)
62                 n, err := slowCopy(wb, bytes.NewBuffer(in), rand.Intn(8192)+1)
63                 c.Check(err, IsNil)
64                 c.Check(n, Equals, int64(insize))
65                 c.Check(wb.Close(), IsNil)
66                 c.Check(pw.Close(), IsNil)
67                 <-done
68                 c.Check(out.Len(), Equals, insize)
69                 for i := 0; i < out.Len() && i < len(in); i++ {
70                         if out.Bytes()[i] != in[i] {
71                                 c.Errorf("content mismatch at byte %d", i)
72                                 break
73                         }
74                 }
75         }
76 }
77
78 func slowCopy(dst io.Writer, src io.Reader, bufsize int) (int64, error) {
79         wrote := int64(0)
80         buf := make([]byte, bufsize)
81         for {
82                 time.Sleep(time.Duration(rand.Intn(100) + 1))
83                 n, err := src.Read(buf)
84                 if n > 0 {
85                         n, err := dst.Write(buf[:n])
86                         wrote += int64(n)
87                         if err != nil {
88                                 return wrote, err
89                         }
90                 }
91                 if err == io.EOF {
92                         return wrote, nil
93                 }
94                 if err != nil {
95                         return wrote, err
96                 }
97         }
98 }