Merge pull request #3 from holiman/less_alloc

all: less memory use on write to disk

Author: holiman

binarymarshaler.go

@@ -14,11 +14,24 @@ import (
 	"bytes"
 	"crypto/sha512"
 	"encoding/binary"
+	"io"
 )
 
+// counter is a utility to count bytes written
+type counter struct {
+	bytes int
+}
+
+func (c *counter) Write(p []byte) (n int, err error) {
+	count := len(p)
+	c.bytes += count
+	return count, nil
+}
+
 // conforms to encoding.BinaryMarshaler
 
-// marshalled binary layout (Little Endian):
+// MarshallToWriter marshalls the filter into the given io.Writer
+// Binary layout (Little Endian):
 //
 //	 k	1 uint64
 //	 n	1 uint64
@@ -29,59 +42,68 @@ import (
 //
 //	 size = (3 + k + (m+63)/64) * 8 bytes
 //
-
-func (f *Filter) marshal() (buf *bytes.Buffer,
-	hash [sha512.Size384]byte,
-	err error,
-) {
+func (f *Filter) MarshallToWriter(out io.Writer) (int, [sha512.Size384]byte, error) {
+	var (
+		c      = &counter{0}
+		hasher = sha512.New384()
+		mw     = io.MultiWriter(out, hasher, c)
+		hash   [sha512.Size384]byte
+	)
 	f.lock.RLock()
 	defer f.lock.RUnlock()
-
 	debug("write bf k=%d n=%d m=%d\n", f.K(), f.n, f.m)
 
-	buf = new(bytes.Buffer)
-
-	err = binary.Write(buf, binary.LittleEndian, f.K())
-	if err != nil {
-		return nil, hash, err
+	if err := binary.Write(mw, binary.LittleEndian, f.K()); err != nil {
+		return c.bytes, hash, err
 	}
-
-	err = binary.Write(buf, binary.LittleEndian, f.n)
-	if err != nil {
-		return nil, hash, err
+	if err := binary.Write(mw, binary.LittleEndian, f.n); err != nil {
+		return c.bytes, hash, err
 	}
-
-	err = binary.Write(buf, binary.LittleEndian, f.m)
-	if err != nil {
-		return nil, hash, err
+	if err := binary.Write(mw, binary.LittleEndian, f.m); err != nil {
+		return c.bytes, hash, err
 	}
-
-	err = binary.Write(buf, binary.LittleEndian, f.keys)
-	if err != nil {
-		return nil, hash, err
+	if err := binary.Write(mw, binary.LittleEndian, f.keys); err != nil {
+		return c.bytes, hash, err
 	}
-
-	err = binary.Write(buf, binary.LittleEndian, f.bits)
+	// Write it in chunks of 5% (but at least 4K). Otherwise, the binary.Write will allocate a
+	// same-size slice of bytes, doubling the memory usage
+	var chunkSize = len(f.bits) / 20
+	if chunkSize < 512 {
+		chunkSize = 512 // Min 4K bytes (512 uint64s)
+	}
+	bs := make([]byte, chunkSize*8)
+	for start := 0; start < len(f.bits); {
+		end := start + chunkSize
+		if end > len(f.bits) {
+			end = len(f.bits)
+		}
+		for i, x := range f.bits[start:end] {
+			binary.LittleEndian.PutUint64(bs[8*i:], x)
+		}
+		if _, err := mw.Write(bs[0 : (end-start)*8]); err != nil {
+			return c.bytes, hash, err
+		}
+		start = end
+	}
+	// Now we stop using the multiwriter, pick out the hash of what we've
+	// written so far, and then write the hash to the output
+	hashbytes := hasher.Sum(nil)
+	copy(hash[:], hashbytes[:sha512.Size384])
+	err := binary.Write(out, binary.LittleEndian, hashbytes)
 	if err != nil {
-		return nil, hash, err
+		debug("bloomfilter.MarshalBinary: Successfully wrote %d byte(s), sha384 %v",
+			c.bytes, hash)
 	}
-
-	hash = sha512.Sum384(buf.Bytes())
-	err = binary.Write(buf, binary.LittleEndian, hash)
-	return buf, hash, err
+	return c.bytes, hash, err
 }
 
 // MarshalBinary converts a Filter into []bytes
 func (f *Filter) MarshalBinary() (data []byte, err error) {
-	buf, hash, err := f.marshal()
+	buf := new(bytes.Buffer)
+	_, _, err = f.MarshallToWriter(buf)
 	if err != nil {
 		return nil, err
 	}
-
-	debug(
-		"bloomfilter.MarshalBinary: Successfully wrote %d byte(s), sha384 %v",
-		buf.Len(), hash,
-	)
 	data = buf.Bytes()
 	return data, nil
 }

fileio.go

@@ -80,12 +80,7 @@ func (f *Filter) WriteTo(w io.Writer) (n int64, err error) {
 		err = rawW.Close()
 	}()
 
-	content, err := f.MarshalBinary()
-	if err != nil {
-		return -1, err
-	}
-
-	intN, err := rawW.Write(content)
+	intN, _, err := f.MarshallToWriter(rawW)
 	n = int64(intN)
 	return n, err
 }

fileio_test.go

@@ -12,6 +12,10 @@ package bloomfilter
 
 import (
 	"bytes"
+	"crypto/sha512"
+	"fmt"
+	"math/rand"
+	"runtime"
 	"testing"
 )
 
@@ -38,3 +42,83 @@ func TestWriteRead(t *testing.T) {
 		t.Error("Filters not equal")
 	}
 }
+
+func bToMb(b uint64) uint64 {
+	return b / 1024 / 1024
+}
+func PrintMemUsage() {
+	var m runtime.MemStats
+	runtime.ReadMemStats(&m)
+	// For info on each, see: https://golang.org/pkg/runtime/#MemStats
+	fmt.Printf("Alloc = %v MiB", bToMb(m.Alloc))
+	fmt.Printf("\tTotalAlloc = %v MiB", bToMb(m.TotalAlloc))
+	fmt.Printf("\tSys = %v MiB", bToMb(m.Sys))
+	fmt.Printf("\tNumGC = %v\n", m.NumGC)
+}
+
+func totAllocMb() uint64 {
+	var m runtime.MemStats
+	runtime.ReadMemStats(&m)
+	return bToMb(m.TotalAlloc)
+}
+
+type devnull struct{}
+func (d devnull) Write(p []byte) (n int, err error) {
+	return len(p), nil
+}
+
+func TestWrite(t *testing.T) {
+	// 1Mb
+	f, _ := New(4*8*1024*1024, 1)
+	fmt.Printf("Allocated 1mb filter\n")
+	PrintMemUsage()
+	_, _ = f.WriteTo(devnull{})
+	fmt.Printf("Wrote filter to devnull\n")
+	PrintMemUsage()
+}
+
+// fillRandom fills the filter with N random values, where N is roughly half
+// the size of the number of uint64's in the filter
+func fillRandom(f *Filter) {
+	num := len(f.bits) * 4
+	for i := 0; i < num; i++ {
+		f.AddHash(uint64(rand.Int63()))
+	}
+}
+
+// TestMarshaller tests that it writes outputs correctly.
+func TestMarshaller(t *testing.T) {
+
+	h1 := sha512.New384()
+	h2 := sha512.New384()
+
+	f, _ := New(1*8*1024*1024, 1)
+	fillRandom(f)
+	// Marshall using writer
+	f.MarshallToWriter(h1)
+	// Marshall as a blob
+	data, _ := f.MarshalBinary()
+	h2.Write(data)
+
+	if have, want := h1.Sum(nil), h2.Sum(nil); !bytes.Equal(have, want) {
+		t.Errorf("Marshalling error, have %x want %x", have, want)
+	}
+}
+
+func BenchmarkWrite1Mb(b *testing.B) {
+
+	// 1Mb
+	f, _ := New(1*8*1024*1024, 1)
+	f.Add(hashableUint64(0))
+	f.Add(hashableUint64(1))
+	f.Add(hashableUint64(1 << 3))
+	f.Add(hashableUint64(1 << 40))
+	f.Add(hashableUint64(1 << 23))
+	f.Add(hashableUint64(1 << 16))
+	f.Add(hashableUint64(1 << 28))
+
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		_, _ = f.WriteTo(devnull{})
+	}
+}

textmarshaler.go

@@ -34,7 +34,7 @@ func (f *Filter) MarshalText() (text []byte, err error) {
 		s += fmt.Sprintf(bitsFormat, w) + nl()
 	}
 
-	_, hash, err := f.marshal()
+	_, hash, err := f.MarshallToWriter(devnull{})
 	if err != nil {
 		return nil, err
 	}

textunmarshaler.go

@@ -78,7 +78,7 @@ func unmarshalAndCheckTextHash(r io.Reader, f *Filter) (err error) {
 		}
 	}
 
-	_, expectedHash, err := f.marshal()
+	_, expectedHash, err := f.MarshallToWriter(devnull{})
 	if err != nil {
 		return err
 	}