// Increment adds a (potentially negative) value to a counter on the server,
// scheduled to be reverted after a certain duration.
func (c *Client) Increment(k uint64, v float64, t time.Duration) error {
return types.WriteMessage(c.s, types.IncrementMessage{
Key: k,
Val: v,
TTE: uint32(time.Now().Add(t).Unix()),
})
}
// ReadOrQuery does what it sounds like it does. It either returns the locally
// cached value of a particular key, or it makes a request to the server.
func (c *Client) ReadOrQuery(k uint64, t time.Duration) (float64, error) {
v, err := c.Read(k)
if err == nil {
return v, nil
}
l := make(chan float64, 1)
c.m.Lock()
c.o[k] = append(c.o[k], l)
c.m.Unlock()
if err := types.WriteMessage(c.s, types.QueryMessage{Key: k}); err != nil {
return 0, err
}
select {
case v := <-l:
return v, nil
case <-time.After(t):
return 0, ErrTimeout
}
}
func main() {
flag.Parse()
level, err := logrus.ParseLevel(*logLevel)
if err != nil {
logrus.WithField("error", err.Error()).Fatal("couldn't parse log level")
}
logrus.SetLevel(level)
var dmtx sync.RWMutex
data := make(map[uint64]float64)
var smtx sync.RWMutex
subs := make(map[uint64][]chan types.Message)
db, err := bolt.Open(*dbFile, 0644, nil)
if err != nil {
panic(err)
}
if err := db.Update(func(tx *bolt.Tx) error {
if _, err := tx.CreateBucketIfNotExists([]byte("deltas")); err != nil {
return err
}
return nil
}); err != nil {
panic(err)
}
if err := db.View(func(tx *bolt.Tx) error {
return tx.Bucket([]byte("deltas")).ForEach(func(kb []byte, vb []byte) error {
var p kv
if err := binary.Read(bytes.NewReader(vb), binary.BigEndian, &p); err != nil {
return err
}
data[p.K] = data[p.K] + p.V
return nil
})
}); err != nil {
panic(err)
}
applyChanges := func() error {
t := uint64(time.Now().Unix())
delta := make(map[uint64]float64)
n := 0
if err := db.Update(func(tx *bolt.Tx) error {
b := tx.Bucket([]byte("deltas"))
c := b.Cursor()
for kb, vb := c.First(); kb != nil; kb, vb = c.Next() {
k := binary.BigEndian.Uint64(kb)
if k > t {
logrus.Debug("found delta for the future; bailing out")
break
}
var p kv
if err := binary.Read(bytes.NewReader(vb), binary.BigEndian, &p); err != nil {
return err
}
delta[p.K] = delta[p.K] + p.V
if err := c.Delete(); err != nil {
return err
}
n++
}
return nil
}); err != nil {
return err
}
logrus.WithField("count", n).Debug("applying deltas")
dmtx.Lock()
defer dmtx.Unlock()
for k, v := range delta {
logrus.WithFields(logrus.Fields{
"key": k,
"value": data[k],
"delta": v,
}).Debug("applying delta")
data[k] = data[k] - v
for _, c := range subs[k] {
c <- types.NotifyMessage{
Key: k,
Val: data[k],
}
}
}
return nil
}
go func() {
for {
if err := applyChanges(); err != nil {
panic(err)
}
time.Sleep(*checkInterval)
}
}()
l, err := net.Listen("tcp4", *addr)
if err != nil {
logrus.WithField("error", err.Error()).Fatal("couldn't open listening socket")
}
logrus.WithField("addr", l.Addr().String()).Info("listening")
for {
s, err := l.Accept()
if err != nil {
if err == io.EOF {
break
}
logrus.WithField("error", err.Error()).Warn("couldn't accept connection")
continue
}
logrus.WithFields(logrus.Fields{
"remote": s.RemoteAddr().String(),
"local": s.LocalAddr().String(),
}).Info("got connection")
go func() {
defer func() {
if e := recover(); e != nil {
logrus.WithFields(logrus.Fields{
"remote": s.RemoteAddr().String(),
"local": s.LocalAddr().String(),
"error": e,
}).Warn("error occurred with client")
}
}()
ch := make(chan types.Message, 10)
mysubs := make(map[uint64]bool)
defer func() {
logrus.WithFields(logrus.Fields{
"remote": s.RemoteAddr().String(),
"local": s.LocalAddr().String(),
}).Info("connection closed")
}()
defer func() {
if len(mysubs) == 0 {
return
}
smtx.Lock()
for k := range mysubs {
subs[k] = remove(subs[k], ch)
}
smtx.Unlock()
close(ch)
}()
go func() {
for m := range ch {
logrus.WithFields(logrus.Fields{
"remote": s.RemoteAddr().String(),
"local": s.LocalAddr().String(),
"type": m.Type().String(),
}).Debug("sending message")
if err := types.WriteMessage(s, m); err != nil {
panic(err)
}
}
}()
b := bufio.NewReader(s)
for {
m, err := types.ReadMessage(b)
if err != nil {
if err == io.EOF {
break
}
panic(err)
}
logrus.WithFields(logrus.Fields{
"remote": s.RemoteAddr().String(),
"local": s.LocalAddr().String(),
"type": m.Type().String(),
}).Debug("received message")
switch m := m.(type) {
case *types.IncrementMessage:
if err := db.Update(func(t *bolt.Tx) error {
k := bytes.NewBuffer(nil)
if err := binary.Write(k, binary.BigEndian, kv{K: uint64(m.TTE), V: rand.Float64()}); err != nil {
return err
}
v := bytes.NewBuffer(nil)
if err := binary.Write(v, binary.BigEndian, kv{K: m.Key, V: m.Val}); err != nil {
return err
}
if err := t.Bucket([]byte("deltas")).Put(k.Bytes(), v.Bytes()); err != nil {
return err
}
return nil
}); err != nil {
panic(err)
}
dmtx.Lock()
data[m.Key] = data[m.Key] + m.Val
for _, c := range subs[m.Key] {
c <- types.NotifyMessage{
Key: m.Key,
Val: data[m.Key],
}
}
dmtx.Unlock()
case *types.SubscribeMessage:
smtx.Lock()
subs[m.Key] = append(subs[m.Key], ch)
mysubs[m.Key] = true
smtx.Unlock()
case *types.UnsubscribeMessage:
smtx.Lock()
subs[m.Key] = remove(subs[m.Key], ch)
delete(mysubs, m.Key)
smtx.Unlock()
case *types.QueryMessage:
dmtx.RLock()
v := data[m.Key]
dmtx.RUnlock()
ch <- types.NotifyMessage{
Key: m.Key,
Val: v,
}
}
}
}()
}
}
// Unsubscribe makes a request to the server to stop sending update
// notifications for a particular key.
func (c *Client) Unsubscribe(k uint64) error {
return types.WriteMessage(c.s, types.UnsubscribeMessage{
Key: k,
})
}