path: root/src/hub/src/spreadspace.org/sfive/s5store.go

//
// sfive
//
// sfive - spreadspace streaming statistics suite is a generic
// statistic collection tool for streaming server infrastuctures.
// The system collects and stores meta data like number of views
// and throughput from a number of streaming servers and stores
// it in a global data store.
// The data acquisition is designed to be generic and extensible in
// order to support different streaming software.
// sfive also contains tools and applications to filter and visualize
// live and recorded data.
//
//
// Copyright (C) 2014-2017 Christian Pointner <equinox@spreadspace.org>
//                         Markus Grüneis <gimpf@gimpf.org>
//
// This file is part of sfive.
//
// sfive is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License version 3
// as published by the Free Software Foundation.
//
// sfive is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with sfive. If not, see <http://www.gnu.org/licenses/>.
//

package sfive

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"time"

	"github.com/coreos/bbolt"
	"github.com/pborman/uuid"
)

const (
	StoreGetUpdatesLimit = 42000
	StoreVersion         = 1
)

var (
	storeBuckets = []string{latestUpdatesBn, hubUUIDsFwdBn, hubUUIDsRevBn, updatesBn,
		sourcesFwdBn, sourcesRevBn, clientsFwdBn, clientsRevBn, clientDataBn, userAgentsFwdBn, userAgentsRevBn}
)

type Store struct {
	version  int64
	hubUUID  string
	db       *bolt.DB
	readOnly bool
}

//
// Initialization and Destruction
//

func openDB(dbPath string, readOnly bool) (db *bolt.DB, version int64, hubUUID string, err error) {
	if _, err = os.Stat(dbPath); err != nil {
		if os.IsNotExist(err) {
			err = nil
		}
		return
	}
	opts := &bolt.Options{Timeout: 100 * time.Millisecond, ReadOnly: readOnly}
	db, err = bolt.Open(dbPath, 0600, opts)
	if err != nil {
		return
	}

	err = db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(hubInfoBn))
		if b == nil {
			return errors.New("store: failed to open, bucket '" + hubInfoBn + "'does not exist.")
		}
		bVersion := b.Get([]byte(storeVersionKey))
		if bVersion == nil {
			return errors.New("store: failed to open, no hub version found.")
		}
		version = btoi(bVersion)
		if version != StoreVersion {
			return fmt.Errorf("store: failed to open, wrong hub version: %d (expected: %d)", version, StoreVersion)
		}

		bHubID := b.Get([]byte(hubUUIDKey))
		if bHubID != nil {
			hubUUID = string(bHubID)
		}
		if hubUUID == "" {
			return errors.New("store: failed to open, UUID does not exist or is empty.")
		}

		for _, bn := range storeBuckets {
			if b := tx.Bucket([]byte(bn)); b == nil {
				return errors.New("store: failed to open, bucket '" + bn + "' does not exist.")
			}
		}
		return nil
	})
	if err != nil {
		db.Close()
	}
	return
}

func createDB(dbPath string) (db *bolt.DB, version int64, hubUUID string, err error) {
	db, err = bolt.Open(dbPath, 0600, &bolt.Options{Timeout: 100 * time.Millisecond})
	if err != nil {
		return
	}

	err = db.Update(func(tx *bolt.Tx) error {
		for _, bn := range storeBuckets {
			if _, err := tx.CreateBucket([]byte(bn)); err != nil {
				return err
			}
		}

		b, err := tx.CreateBucket([]byte(hubInfoBn))
		if err != nil {
			return err
		}
		version = int64(StoreVersion)
		if err := b.Put([]byte(storeVersionKey), itob(version)); err != nil {
			return err
		}
		hubUUID = uuid.New()
		if err := b.Put([]byte(hubUUIDKey), []byte(hubUUID)); err != nil {
			return err
		}
		return nil
	})
	if err != nil {
		db.Close()
	}

	return
}

func NewStore(cfg StoreConfig) (*Store, error) {
	db, version, hubid, err := openDB(cfg.DBPath, cfg.ReadOnly)
	if err != nil {
		return nil, err
	}

	if db != nil {
		if cfg.ReadOnly {
			s5l.Printf("store: opened read-only (UUID: %s)", hubid)
		} else {
			s5l.Printf("store: opened (UUID: %s)", hubid)
		}
		return &Store{version, hubid, db, cfg.ReadOnly}, nil
	}

	if cfg.ReadOnly {
		return nil, errors.New("store: failed to open, requested read-only mode but store file does not exist.")
	}

	db, version, hubid, err = createDB(cfg.DBPath)
	if err != nil {
		return nil, err
	}
	s5l.Printf("store: initialized (UUID: %s)", hubid)
	return &Store{version, hubid, db, cfg.ReadOnly}, nil
}

func (st *Store) Close() {
	s5l.Printf("store: closing")
	st.db.Close()
}

//
// Append data
//

// append key-value pairs to buckets

func (st *Store) insertNewHub(tx *bolt.Tx, hubUUID string) (hubID int64, err error) {
	if hubUUID == "" {
		return
	}

	bf := tx.Bucket([]byte(hubUUIDsFwdBn))
	bf.FillPercent = 1.0 // we only do appends
	br := tx.Bucket([]byte(hubUUIDsRevBn))
	br.FillPercent = 1.0 // we only do appends

	bHubID := bf.Get([]byte(hubUUID))
	if bHubID != nil {
		return btoi(bHubID), nil
	}

	next, _ := bf.NextSequence()
	hubID = int64(next)
	if err = bf.Put([]byte(hubUUID), itob(hubID)); err != nil {
		return
	}
	if err = br.Put(itob(hubID), []byte(hubUUID)); err != nil {
		return
	}

	return hubID, err
}

func (st *Store) insertNewSource(tx *bolt.Tx, src *sourceDB) (srcID int64, err error) {
	bf := tx.Bucket([]byte(sourcesFwdBn))
	// br.FillPercent = 1.0 // these appends are not ordered (the key is the slug and not an integer id)
	br := tx.Bucket([]byte(sourcesRevBn))
	br.FillPercent = 1.0 // we only do appends (with ever incrementing interger ids)

	slug := src.Slug()
	bSrcID := bf.Get([]byte(slug))
	if bSrcID != nil {
		return btoi(bSrcID), nil
	}

	var jsonData []byte
	if jsonData, err = json.Marshal(src); err != nil {
		return
	}

	next, _ := bf.NextSequence()
	srcID = int64(next)
	if err = bf.Put([]byte(slug), itob(srcID)); err != nil {
		return
	}
	if err = br.Put(itob(srcID), jsonData); err != nil {
		return
	}

	return srcID, err
}

func (st *Store) insertUpdate(tx *bolt.Tx, du *updateDB) (duID int64, err error) {
	b := tx.Bucket([]byte(updatesBn))
	b.FillPercent = 1.0 // we only do appends

	next, _ := b.NextSequence()
	duID = int64(next)

	var jsonData []byte
	if jsonData, err = json.Marshal(du); err != nil {
		return
	}
	err = b.Put(itob(duID), jsonData)
	return
}

func (st *Store) insertNewUserAgent(tx *bolt.Tx, ua string) (uaID int64, err error) {
	bf := tx.Bucket([]byte(userAgentsFwdBn))
	bf.FillPercent = 1.0 // we only do appends
	br := tx.Bucket([]byte(userAgentsRevBn))
	br.FillPercent = 1.0 // we only do appends

	bUaID := bf.Get([]byte(ua))
	if bUaID != nil {
		return btoi(bUaID), nil
	}

	next, _ := bf.NextSequence()
	uaID = int64(next)
	if err = bf.Put([]byte(ua), itob(uaID)); err != nil {
		return
	}
	if err = br.Put(itob(uaID), []byte(ua)); err != nil {
		return
	}

	return uaID, err
}

func (st *Store) insertNewClient(tx *bolt.Tx, client *clientDB) (clientID int64, err error) {
	bf := tx.Bucket([]byte(clientsFwdBn))
	// br.FillPercent = 1.0 // these appends are not ordered (the key is the slug and not an integer id)
	br := tx.Bucket([]byte(clientsRevBn))
	br.FillPercent = 1.0 // we only do appends (with ever incrementing interger ids)

	slug := client.Slug()
	bClientID := bf.Get([]byte(slug))
	if bClientID != nil {
		return btoi(bClientID), nil
	}

	var jsonData []byte
	if jsonData, err = json.Marshal(client); err != nil {
		return
	}

	next, _ := bf.NextSequence()
	clientID = int64(next)
	if err = bf.Put([]byte(slug), itob(clientID)); err != nil {
		return
	}
	if err = br.Put(itob(clientID), jsonData); err != nil {
		return
	}

	return clientID, err
}

func (st *Store) insertClients(tx *bolt.Tx, uID int64, cd []Client) error {
	if len(cd) == 0 {
		return nil
	}

	b := tx.Bucket([]byte(clientDataBn))
	b.FillPercent = 1.0 // we only do appends

	data := []clientDataDB{}
	for _, c := range cd {
		uaID, err := st.insertNewUserAgent(tx, c.UserAgent)
		if err != nil {
			return err
		}

		clientID, err := st.insertNewClient(tx, NewClientDB(&c))
		if err != nil {
			return err
		}

		data = append(data, clientDataDB{clientID, uaID, c.BytesSent})
	}

	jsonData, err := json.Marshal(data)
	if err != nil {
		return err
	}
	return b.Put(itob(uID), jsonData)
}

func (st *Store) setLastUpdateForUUID(tx *bolt.Tx, uuid string, uID int64) error {
	b := tx.Bucket([]byte(latestUpdatesBn))
	b.FillPercent = 1.0 // we only do appends

	last := b.Get([]byte(uuid))
	if last != nil && btoi(last) > uID {
		return nil
	}
	return b.Put([]byte(uuid), itob(uID))
}

// Split up the multidimensional dataupdate and append all the key-value pairs

func (st *Store) appendItem(tx *bolt.Tx, uf *UpdateFull) (uID int64, err error) {
	u := NewUpdateDB(uf)
	s := NewSourceDB(uf)

	if u.SourceHubID, err = st.insertNewHub(tx, uf.SourceHubUUID); err != nil {
		return
	}
	if u.SourceID, err = st.insertNewSource(tx, s); err != nil {
		return
	}
	if uID, err = st.insertUpdate(tx, u); err != nil {
		return
	}
	if err = st.insertClients(tx, uID, uf.Data.Clients); err != nil {
		return
	}

	if uf.SourceHubUUID != "" {
		err = st.setLastUpdateForUUID(tx, uf.SourceHubUUID, u.SourceHubUpdateID)
	}
	if uf.ForwardHubUUID != "" {
		err = st.setLastUpdateForUUID(tx, uf.ForwardHubUUID, uf.ForwardHubUpdateID)
	}
	return
}

// Public Append Interface

func (st *Store) AppendMany(updates []*UpdateFull) (err error) {
	if st.readOnly {
		return ErrReadOnly
	}
	return st.db.Update(func(tx *bolt.Tx) error {
		for _, update := range updates {
			if _, err := st.appendItem(tx, update); err != nil {
				return err
			}
		}
		return nil
	})
}

func (st *Store) Append(update *UpdateFull) (err error) {
	return st.AppendMany([]*UpdateFull{update})
}

//
// Fetch data
//

// fetch key-value pairs from buckets

func (st *Store) getHub(tx *bolt.Tx, id int64) string {
	b := tx.Bucket([]byte(hubUUIDsRevBn))
	uuid := b.Get(itob(id))
	if uuid == nil {
		return ""
	}
	return string(uuid)
}

func (st *Store) getSource(tx *bolt.Tx, id int64) (source *sourceDB, err error) {
	b := tx.Bucket([]byte(sourcesRevBn))

	jsonData := b.Get(itob(id))
	if jsonData == nil {
		err = ErrSourceNotFound
		return
	}
	source = &sourceDB{}
	if err = json.Unmarshal(jsonData, source); err != nil {
		return
	}
	return
}

func (st *Store) getClients(tx *bolt.Tx, id int64) (clients []Client, err error) {
	bcd := tx.Bucket([]byte(clientDataBn))
	bc := tx.Bucket([]byte(clientsRevBn))
	bu := tx.Bucket([]byte(userAgentsRevBn))

	jsonData := bcd.Get(itob(id))
	if jsonData == nil {
		return
	}
	data := []clientDataDB{}
	if err = json.Unmarshal(jsonData, &data); err != nil {
		return
	}
	for _, c := range data {
		cd := Client{BytesSent: c.BytesSent}
		jsonData := bc.Get(itob(c.ClientID))
		if jsonData == nil {
			err = ErrClientNotFound
			return
		}
		client := &clientDB{}
		if err = json.Unmarshal(jsonData, client); err != nil {
			return
		}
		cd.CopyFromClientDB(client)

		ua := bu.Get(itob(c.UserAgentID))
		if ua != nil {
			cd.UserAgent = string(ua)
		}
		clients = append(clients, cd)
	}
	return
}

// fetch all the key-value pairs and merge them into the multidimensional dataupdate

func (st *Store) fetchItem(tx *bolt.Tx, uID int64, u *updateDB) (update *UpdateFull, err error) {
	update = &UpdateFull{}
	update.CopyFromUpdateDB(u, st.getHub(tx, u.SourceHubID), st.hubUUID, uID)
	var src *sourceDB
	if src, err = st.getSource(tx, u.SourceID); err != nil {
		return
	}
	update.CopyFromSourceDB(src)
	if update.Data.Clients, err = st.getClients(tx, uID); err != nil {
		return
	}
	return
}

// Public Fetch Interface

func (st *Store) GetUpdatesAfter(id, limit int64) (updates []*UpdateFull, err error) {
	updates = []*UpdateFull{}
	if id < 0 { // TODO: interpret negative values as last x values
		id = 0
	}
	if limit < 0 || limit > StoreGetUpdatesLimit {
		if limit > 0 {
			s5l.Printf("store: truncating get-update limit to %d (from %d)", StoreGetUpdatesLimit, limit)
		}
		limit = StoreGetUpdatesLimit
	}
	err = st.db.View(func(tx *bolt.Tx) error {
		c := tx.Bucket([]byte(updatesBn)).Cursor()
		k, v := c.Seek(itob(id))
		if k == nil {
			return nil
		}
		if btoi(k) == id {
			k, v = c.Next()
		}
		for ; k != nil; k, v = c.Next() {
			var d *updateDB
			if err := json.Unmarshal(v, &d); err != nil {
				return err
			}

			duf, err := st.fetchItem(tx, btoi(k), d)
			if err != nil {
				return err
			}
			updates = append(updates, duf)
			if int64(len(updates)) >= limit {
				return nil
			}
		}
		return nil
	})
	return
}

func (st *Store) GetUpdate(id int64) (update *UpdateFull, err error) {
	err = st.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(updatesBn))

		jsonData := b.Get(itob(id))
		if jsonData == nil {
			return ErrNotFound
		}

		var d *updateDB
		if err := json.Unmarshal(jsonData, &d); err != nil {
			return err
		}

		var err error
		if update, err = st.fetchItem(tx, id, d); err != nil {
			return err
		}
		return nil
	})
	return
}

//
// Auxilliary Data
//

func (st *Store) GetHubUUID() string {
	return st.hubUUID
}

func (st *Store) GetLastUpdateID() (updateID int64, err error) {
	err = st.db.View(func(tx *bolt.Tx) error {
		updateID = int64(tx.Bucket([]byte(updatesBn)).Sequence())
		return nil
	})
	return
}

func (st *Store) GetLastUpdateIDForUUID(uuid string) (updateID int64, err error) {
	if uuid == st.hubUUID {
		return st.GetLastUpdateID()
	}

	err = st.db.View(func(tx *bolt.Tx) error {
		bUpdateID := tx.Bucket([]byte(latestUpdatesBn)).Get([]byte(uuid))
		if bUpdateID == nil {
			return nil
		}
		updateID = btoi(bUpdateID)
		return nil
	})
	return
}

func (st *Store) GetHubs() (hubs []string, err error) {
	hubs = []string{st.hubUUID}
	err = st.db.View(func(tx *bolt.Tx) error {
		c := tx.Bucket([]byte(hubUUIDsRevBn)).Cursor()
		for k, v := c.First(); k != nil; k, v = c.Next() {
			hubs = append(hubs, string(v))
		}
		return nil
	})
	return
}

func (st *Store) GetSources() (sources []*Source, err error) {
	sources = []*Source{}
	err = st.db.View(func(tx *bolt.Tx) error {
		c := tx.Bucket([]byte(sourcesRevBn)).Cursor()
		for k, v := c.First(); k != nil; k, v = c.Next() {
			var s *sourceDB
			if err := json.Unmarshal(v, &s); err != nil {
				return err
			}
			src := &Source{}
			src.CopyFromSourceDB(s)
			sources = append(sources, src)
		}
		return nil
	})
	return
}

func (st *Store) GetClients() (clients []*Client, err error) {
	clients = []*Client{}
	err = st.db.View(func(tx *bolt.Tx) error {
		c := tx.Bucket([]byte(clientsRevBn)).Cursor()
		for k, v := c.First(); k != nil; k, v = c.Next() {
			var s *clientDB
			if err := json.Unmarshal(v, &s); err != nil {
				return err
			}
			client := &Client{}
			client.CopyFromClientDB(s)
			clients = append(clients, client)
		}
		return nil
	})
	return
}