Skip to content

Commit

Permalink
feat(countrw): add countrw package
Browse files Browse the repository at this point in the history
  • Loading branch information
sentriz committed Mar 10, 2022
1 parent c786276 commit 5155dee
Show file tree
Hide file tree
Showing 2 changed files with 52 additions and 0 deletions.
26 changes: 26 additions & 0 deletions countrw/cr.go
@@ -0,0 +1,26 @@
package countrw

import (
"io"
"sync/atomic"
)

type CountReader struct {
r io.Reader
c *uint64
}

func NewCountReader(r io.Reader) *CountReader {
return &CountReader{r: r, c: new(uint64)}
}

func (c *CountReader) Reset() { atomic.StoreUint64(c.c, 0) }
func (c *CountReader) Count() uint64 { return atomic.LoadUint64(c.c) }

func (c *CountReader) Read(p []byte) (int, error) {
n, err := c.r.Read(p)
atomic.AddUint64(c.c, uint64(n))
return n, err
}

var _ io.Reader = (*CountReader)(nil)
26 changes: 26 additions & 0 deletions countrw/cw.go
@@ -0,0 +1,26 @@
package countrw

import (
"io"
"sync/atomic"
)

type CountWriter struct {
r io.Writer
c *uint64
}

func NewCountWriter(r io.Writer) *CountWriter {
return &CountWriter{r: r, c: new(uint64)}
}

func (c *CountWriter) Reset() { atomic.StoreUint64(c.c, 0) }
func (c *CountWriter) Count() uint64 { return atomic.LoadUint64(c.c) }

func (c *CountWriter) Write(p []byte) (int, error) {
n, err := c.r.Write(p)
atomic.AddUint64(c.c, uint64(n))
return n, err
}

var _ io.Writer = (*CountWriter)(nil)

0 comments on commit 5155dee

Please sign in to comment.