Add distinct IP counter

common/ipsetsink/sinkcluster/common.go

@@ -0,0 +1,9 @@
+package sinkcluster
+
+import "time"
+
+type SinkEntry struct {
+	RecordingStart time.Time `json:"recordingStart"`
+	RecordingEnd   time.Time `json:"recordingEnd"`
+	Recorded       []byte    `json:"recorded"`
+}

common/ipsetsink/sinkcluster/reader.go

@@ -0,0 +1,60 @@
+package sinkcluster
+
+import (
+	"bufio"
+	"encoding/json"
+	"github.com/clarkduvall/hyperloglog"
+	"io"
+	"time"
+)
+
+func NewClusterCounter(from time.Time, to time.Time) *ClusterCounter {
+	return &ClusterCounter{from: from, to: to}
+}
+
+type ClusterCounter struct {
+	from time.Time
+	to   time.Time
+}
+
+type ClusterCountResult struct {
+	Sum           uint64
+	ChunkIncluded int64
+}
+
+func (c ClusterCounter) Count(reader io.Reader) (*ClusterCountResult, error) {
+	result := ClusterCountResult{}
+	counter, err := hyperloglog.NewPlus(18)
+	if err != nil {
+		return nil, err
+	}
+	inputScanner := bufio.NewScanner(reader)
+	for inputScanner.Scan() {
+		inputLine := inputScanner.Bytes()
+		sinkInfo := SinkEntry{}
+		if err := json.Unmarshal(inputLine, &sinkInfo); err != nil {
+			return nil, err
+		}
+
+		if (sinkInfo.RecordingStart.Before(c.from) && !sinkInfo.RecordingStart.Equal(c.from)) ||
+			sinkInfo.RecordingEnd.After(c.to) {
+			continue
+		}
+
+		restoredCounter, err := hyperloglog.NewPlus(18)
+		if err != nil {
+			return nil, err
+		}
+		err = restoredCounter.GobDecode(sinkInfo.Recorded)
+		if err != nil {
+			return nil, err
+		}
+		result.ChunkIncluded++
+		err = counter.Merge(restoredCounter)
+		if err != nil {
+			return nil, err
+		}
+	}
+	result.Sum = counter.Count()
+	return &result, nil
+}

common/ipsetsink/sinkcluster/writer.go

@@ -0,0 +1,68 @@
+package sinkcluster
+
+import (
+	"bytes"
+	"encoding/json"
+	"io"
+	"log"
+	"time"
+
+	"git.torproject.org/pluggable-transports/snowflake.git/v2/common/ipsetsink"
+)
+
+func NewClusterWriter(writer WriteSyncer, writeInterval time.Duration, sink *ipsetsink.IPSetSink) *ClusterWriter {
+	c := &ClusterWriter{
+		writer:        writer,
+		lastWriteTime: time.Now(),
+		writeInterval: writeInterval,
+		current:       sink,
+	}
+	return c
+}
+
+type ClusterWriter struct {
+	writer        WriteSyncer
+	lastWriteTime time.Time
+	writeInterval time.Duration
+	current       *ipsetsink.IPSetSink
+}
+
+type WriteSyncer interface {
+	Sync() error
+	io.Writer
+}
+
+func (c *ClusterWriter) WriteIPSetToDisk() {
+	currentTime := time.Now()
+	data, err := c.current.Dump()
+	if err != nil {
+		log.Println("unable able to write ipset to file:", err)
+		return
+	}
+	entry := &SinkEntry{
+		RecordingStart: c.lastWriteTime,
+		RecordingEnd:   currentTime,
+		Recorded:       data,
+	}
+	jsonData, err := json.Marshal(entry)
+	if err != nil {
+		log.Println("unable able to write ipset to file:", err)
+		return
+	}
+	jsonData = append(jsonData, byte('\n'))
+	_, err = io.Copy(c.writer, bytes.NewReader(jsonData))
+	if err != nil {
+		log.Println("unable able to write ipset to file:", err)
+		return
+	}
+	c.writer.Sync()
+	c.lastWriteTime = currentTime
+	c.current.Reset()
+}
+
+func (c *ClusterWriter) AddIPToSet(ipAddress string) {
+	if c.lastWriteTime.Add(c.writeInterval).Before(time.Now()) {
+		c.WriteIPSetToDisk()
+	}
+	c.current.AddIPToSet(ipAddress)
+}

common/ipsetsink/sinkcluster/writer_test.go

@@ -0,0 +1,33 @@
+package sinkcluster
+
+import (
+	"bytes"
+	"io"
+	"testing"
+	"time"
+
+	"git.torproject.org/pluggable-transports/snowflake.git/v2/common/ipsetsink"
+
+	. "github.com/smartystreets/goconvey/convey"
+)
+
+type writerStub struct {
+	io.Writer
+}
+
+func (w writerStub) Sync() error {
+	return nil
+}
+
+func TestSinkWriter(t *testing.T) {
+
+	Convey("Context", t, func() {
+		buffer := bytes.NewBuffer(nil)
+		writerStubInst := &writerStub{buffer}
+		sink := ipsetsink.NewIPSetSink("demo")
+		clusterWriter := NewClusterWriter(writerStubInst, time.Minute, sink)
+		clusterWriter.AddIPToSet("1")
+		clusterWriter.WriteIPSetToDisk()
+		So(buffer.Bytes(), ShouldNotBeNil)
+	})
+}