func (c *Client) read() { for { m, err := types.ReadMessage(c.s) if err != nil { if err == io.EOF { break } panic(err) } switch m := m.(type) { case *types.NotifyMessage: c.m.Lock() c.d[m.Key] = m.Val l, ok := c.o[m.Key] if ok { delete(c.o, m.Key) } c.m.Unlock() for _, c := range l { c <- m.Val } for _, ch := range c.c { ch <- KV{K: m.Key, V: m.Val} } } } }
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, } } } }() } }