diff options
author | Jacob Vosmaer <jacob@gitlab.com> | 2019-01-24 19:30:22 +0300 |
---|---|---|
committer | John Cai <jcai@gitlab.com> | 2019-01-24 19:30:22 +0300 |
commit | f1b09f21149be2fb42930dc0730cb8fc786e54af (patch) | |
tree | 492928bd22f3a81cd5bc2252a0fff92d964f3259 /internal/helper/chunk | |
parent | 6cdf9a73866edc5aff7c1f15f059a156948821d2 (diff) |
Refactor refnames RPC's to use chunker
Diffstat (limited to 'internal/helper/chunk')
-rw-r--r-- | internal/helper/chunk/chunker.go | 59 |
1 files changed, 59 insertions, 0 deletions
diff --git a/internal/helper/chunk/chunker.go b/internal/helper/chunk/chunker.go new file mode 100644 index 000000000..48d99f03f --- /dev/null +++ b/internal/helper/chunk/chunker.go @@ -0,0 +1,59 @@ +package chunk + +// Item could be e.g. a commit in an RPC that returns a chunked stream of +// commits. +type Item interface{} + +// Sender encapsulates a gRPC response stream and the current chunk +// that's being built. +// +// Reset, Append, [Append...], Send, Reset, Append, [Append...], Send, ... +type Sender interface { + // Reset should create a fresh response message. + Reset() + // Append should append the given item to the slice in the current response message + Append(Item) + // Send should send the current response message + Send() error +} + +// New returns a new Chunker. +func New(s Sender) *Chunker { return &Chunker{s: s} } + +// Chunker lets you spread items you want to send over multiple chunks. +// This type is not thread-safe. +type Chunker struct { + s Sender + n int +} + +// Send will append an item to the current chunk and send the chunk if it is full. +func (c *Chunker) Send(it Item) error { + if c.n == 0 { + c.s.Reset() + } + + c.s.Append(it) + c.n++ + + const chunkSize = 20 + if c.n >= chunkSize { + return c.sendResponseMsg() + } + + return nil +} + +func (c *Chunker) sendResponseMsg() error { + c.n = 0 + return c.s.Send() +} + +// Flush sends remaining items in the current chunk, if any. +func (c *Chunker) Flush() error { + if c.n == 0 { + return nil + } + + return c.sendResponseMsg() +} |