예제 #1
0
파일: client.go 프로젝트: deoxxa/counterd
// 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()),
	})
}
예제 #2
0
파일: client.go 프로젝트: deoxxa/counterd
// 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
	}
}
예제 #3
0
파일: main.go 프로젝트: deoxxa/counterd
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,
					}
				}
			}
		}()
	}
}
예제 #4
0
파일: client.go 프로젝트: deoxxa/counterd
// 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,
	})
}