funmow/db.go

package funmow

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/boltdb/bolt"
	"strconv"
)

const (
	DEBUG = false
)

type DBRef int

func NewDBRefFromHashRef(v string) (DBRef, error) {
	var destInt int
	n, err := fmt.Sscanf(v, "#%d", &destInt)

	if err != nil || n == 0 {
		return 0, err
	}

	return DBRef(destInt), nil
}

func NewDBRefFromString(v string) (DBRef, error) {
	intVal, err := strconv.Atoi(v)
	return DBRef(intVal), err
}

func NewDBRefFromKey(b []byte) DBRef {
	return DBRef(binary.BigEndian.Uint64(b))
}

func NewDBRefFromChildKey(b []byte) DBRef { // returns the child part of a childkey
	return DBRef(binary.BigEndian.Uint64(b[8:]))
}

func (r DBRef) Valid() bool {
	return r != 0
}

func (r DBRef) Limbo() bool {
	return r == 0
}

func (r DBRef) String() string {
	return strconv.Itoa(int(r))
}

func (r DBRef) Key() []byte {
	b := make([]byte, 8)
	binary.BigEndian.PutUint64(b, uint64(r))
	return b
}

func (r DBRef) ChildKey(s DBRef) []byte {
	b := make([]byte, 16)
	binary.BigEndian.PutUint64(b, uint64(r))
	binary.BigEndian.PutUint64(b[8:], uint64(s))
	return b
}

type DBRefList []DBRef

func (l DBRefList) First() DBRef {
	if len(l) > 0 {
		return l[0]
	} else {
		return 0
	}
}

type DB struct {
	path string
	db   *bolt.DB
}

func NewDB(path string) *DB {
	return &DB{path: path}
}

func (s *DB) Close() {
	s.db.Close()
}

func (s *DB) Open() error {
	var err error
	s.db, err = bolt.Open(s.path, 0600, nil)
	if err != nil {
		return err
	}
	return s.db.Update(func(tx *bolt.Tx) error {
		var err error
		objectBucket, err := tx.CreateBucketIfNotExists([]byte("object"))
		if err != nil {
			return fmt.Errorf("create bucket: %s", err)
		}

		_, err = tx.CreateBucketIfNotExists([]byte("child"))
		if err != nil {
			return fmt.Errorf("create bucket: %s", err)
		}

		_, err = tx.CreateBucketIfNotExists([]byte("parent"))
		if err != nil {
			return fmt.Errorf("create bucket: %s", err)
		}

		_, err = tx.CreateBucketIfNotExists([]byte("player"))
		if err != nil {
			return fmt.Errorf("create bucket: %s", err)
		}

		// make sure we have a limbo
		if _, found := s.txRetrieveObject(objectBucket, 0); !found {
			limbo := Object{Name: "Limbo", Description: "There's very little to see here.", ID: 0}
			err = s.txStoreObject(objectBucket, limbo, 0)
			if err != nil {
				return err
			}
		}

		return nil
	})
}

func (s *DB) Allocate() (DBRef, error) {
	var id DBRef
	err := s.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("object"))
		seq, err := b.NextSequence()
		if err != nil {
			return err
		}
		id = DBRef(seq)
		return s.txStoreObject(b, Object{ID: id}, id)
	})
	return id, err
}

func (s *DB) Fetch(r DBRef) (Object, bool) { // this has become simply an alias for RetrieveObject

	return s.RetrieveObject(r)
}

func (s *DB) DumpObject(r DBRef) string {
	var dump string
	s.db.View(func(tx *bolt.Tx) error {
		objectBucket := tx.Bucket([]byte("object"))
		childBucket := tx.Bucket([]byte("child"))
		parentBucket := tx.Bucket([]byte("parent"))

		objectDump := string(objectBucket.Get(r.Key()))

		intermediate := make([]string, 0)
		children := s.txGetChildren(childBucket, r)
		for childID, childType := range children {
			intermediate = append(intermediate, fmt.Sprintf("#%d (%s)", childID, childType))
		}

		parent, hasParent := s.txGetParent(parentBucket, r)

		dump = fmt.Sprintf("Object: %s\nHas Parent: %t\nParent ID: #%d\nChildren: %s", objectDump, hasParent, parent, intermediate)

		return nil
	})
	return dump
}

func (s *DB) Delete(objectID DBRef) error {
	return s.db.Update(func(tx *bolt.Tx) error {

		objectBucket := tx.Bucket([]byte("object"))
		childBucket := tx.Bucket([]byte("child"))
		parentBucket := tx.Bucket([]byte("parent"))

		_, found := s.txRetrieveObject(objectBucket, objectID)

		if !found {
			return nil
		}

		parentID, parentFound := s.txGetParent(parentBucket, objectID)

		if parentFound {
			//fmt.Printf("Unlinking from parent #%d\n", parentID)
			err := s.txUnlink(childBucket, parentBucket, parentID, objectID)
			if err != nil {
				return err
			}
		}

		children := s.txGetChildren(childBucket, objectID)

		for childID, childType := range children {
			//fmt.Printf("Unlinking child #%d\n",childID)
			err := s.txUnlink(childBucket, parentBucket, objectID, childID)
			if err != nil {
				return err
			}
			if childType == "exit" {
				//fmt.Printf("Deleting exit child #%d\n", childID)
				err := s.txDelete(objectBucket, childID) // this is bad and will create orphans
				if err != nil {
					return err
				}
			} else {
				// even if parentfound == false, parentID will = 0 (the default value for rooms)
				// that will cause the contents of destroyed rooms to end up in limbo, instead of as orpans
				//fmt.Printf("Linking child #%d to parent #%d\n",childID, parentID)
				err := s.txLink(childBucket, parentBucket, parentID, childID, childType)
				if err != nil {
					return err
				}
			}
		}
		//fmt.Printf("Deleting object #%d\n", objectID)
		err := s.txDelete(objectBucket, objectID)
		if err != nil {
			return err
		}

		if DEBUG {
			fmt.Printf("Delete name: id: %d\n", objectID)
		}

		return nil
	})
}

func (s *DB) SetPlayerID(name string, id DBRef) error {

	err := s.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("player"))
		if b == nil {
			return errors.New("Player bucket not found")
		}
		return b.Put([]byte(name), id.Key())
	})

	if DEBUG {
		fmt.Printf("SetPlayerID name: %s id: %d\n", name, id)
	}

	return err

}

func (s *DB) GetPlayerID(name string) (DBRef, error) {
	var id DBRef
	err := s.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("player"))
		if b == nil {
			return errors.New("Player bucket not found")
		}
		v := b.Get([]byte(name))
		if v == nil {
			return errors.New("Player not found")
		}
		id = NewDBRefFromKey(v)
		return nil
	})

	if DEBUG {
		fmt.Printf("GetPlayerID name: %s id: %d\n", name, id)
	}

	return id, err
}

func (s *DB) CreatePlayer(name string) (DBRef, error) {

	var playerID DBRef

	err := s.db.Update(func(tx *bolt.Tx) error {
		objectBucket := tx.Bucket([]byte("object"))
		playerBucket := tx.Bucket([]byte("player"))
		seq, err := objectBucket.NextSequence()
		if err != nil {
			return err
		}
		playerID = DBRef(seq)
		s.txStoreObject(objectBucket, Object{ID: playerID, Name: name, Type: "player"}, playerID)
		err = playerBucket.Put([]byte(name), playerID.Key())
		return err
	})

	if DEBUG {
		fmt.Printf("CreatePlayer name: %s id: %d\n", name, playerID)
	}

	return playerID, err
}

// All of these functions exist in two forms; one for use inside a transaction and one
// for use outside a transaction.

func (s *DB) RetrieveObject(id DBRef) (Object, bool) {

	o := Object{}
	f := false
	s.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("object"))
		o, f = s.txRetrieveObject(b, id)
		return nil
	})
	return o, f

}

func (s *DB) txRetrieveObject(b *bolt.Bucket, id DBRef) (Object, bool) {

	o := Object{}
	f := false
	v := b.Get(id.Key())
	if v != nil {
		err := json.Unmarshal(v, &o)
		if err == nil {
			f = true

			if DEBUG {
				fmt.Printf("txRetrieveObject id: %d\n", id)
			}

		}
	}
	o.db = s // bit ugly but saves a lot of headaches

	return o, f

}

func (s *DB) StoreObject(o Object, id DBRef) error {

	return s.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("object"))
		return s.txStoreObject(b, o, id)
	})

}

func (s *DB) txStoreObject(b *bolt.Bucket, o Object, id DBRef) error {

	buf, err := json.Marshal(o)

	//fmt.Println("txStoreObject", id.String(), string(buf))

	if err != nil {
		return err
	}
	err = b.Put(id.Key(), buf)
	if err != nil {
		return err
	}

	if DEBUG {
		fmt.Printf("txStoreObject id: %d\n", id)
	}

	return nil

}

func (s *DB) GetChildren(src DBRef) map[DBRef]string {

	var l map[DBRef]string
	s.db.View(func(tx *bolt.Tx) error {
		c := tx.Bucket([]byte("child"))
		l = s.txGetChildren(c, src)
		return nil
	})
	return l

}

func (s *DB) txGetChildren(b *bolt.Bucket, src DBRef) map[DBRef]string {

	l := make(map[DBRef]string)
	prefix := src.Key()      // 00001234
	start := src.ChildKey(0) // 0000123400000000
	c := b.Cursor()
	for k, v := c.Seek(start); bytes.HasPrefix(k, prefix); k, v = c.Next() {
		l[NewDBRefFromChildKey(k)] = string(v)
	}
	return l

}

func (s *DB) GetParent(src DBRef) (DBRef, bool) {

	p := DBRef(0)
	f := false

	s.db.View(func(tx *bolt.Tx) error {
		parentBucket := tx.Bucket([]byte("parent"))
		p, f = s.txGetParent(parentBucket, src)
		return nil
	})
	return p, f

}

func (s *DB) txGetParent(parentBucket *bolt.Bucket, childID DBRef) (DBRef, bool) {

	parentID := DBRef(0)
	f := false

	v := parentBucket.Get(childID.Key())
	if v != nil {
		parentID = NewDBRefFromKey(v)
		f = true
	}

	if DEBUG {
		fmt.Printf("txGetParent child: %d parent: %d\n", childID, parentID)
	}

	//fmt.Printf("txGetParent src: %d found: %t parent: %d\n", src, f, p)
	return parentID, f

}

func (s *DB) IsParent(src DBRef, dest DBRef) bool {

	f := false
	s.db.View(func(tx *bolt.Tx) error {
		parentBucket := tx.Bucket([]byte("parent"))
		f = s.txIsParent(parentBucket, src, dest)
		return nil
	})
	return f

}

func (s *DB) txIsParent(b *bolt.Bucket, src DBRef, dst DBRef) bool {

	v := b.Get(src.Key())
	if v != nil {
		return (bytes.Compare(v, dst.Key()) == 0)
	}
	return false

}

func (s *DB) IsChild(src DBRef, dest DBRef) (string, bool) {

	t := ""
	f := false
	s.db.View(func(tx *bolt.Tx) error {
		childBucket := tx.Bucket([]byte("child"))
		t, f = s.txIsChild(childBucket, src, dest)
		return nil
	})
	return t, f

}

func (s *DB) txIsChild(b *bolt.Bucket, src DBRef, dest DBRef) (string, bool) {

	v := b.Get(src.ChildKey(dest))
	if v != nil {
		return string(v), true
	}
	return "", false

}

func (s *DB) Link(parentID DBRef, childID DBRef, linkType string) error {

	return s.db.Update(func(tx *bolt.Tx) error {
		childBucket := tx.Bucket([]byte("child"))
		parentBucket := tx.Bucket([]byte("parent"))

		return s.txLink(childBucket, parentBucket, parentID, childID, linkType)
	})

}

func (s *DB) txLink(childBucket *bolt.Bucket, parentBucket *bolt.Bucket, parentID DBRef, childID DBRef, linkType string) error {

	//fmt.Printf("txLink attempting to put %d inside %d\n", childID, parentID)

	prevParentID, found := s.txGetParent(parentBucket, childID)
	if found {
		//fmt.Printf("txLink removing %d from previous parent %d\n", childID, prevParentID)
		err := s.txUnlink(childBucket, parentBucket, prevParentID, childID)
		if err != nil {
			return err
		}
	}

	err := childBucket.Put(parentID.ChildKey(childID), []byte(linkType))
	if err != nil {
		return err
	}

	err = parentBucket.Put(childID.Key(), parentID.Key())
	if err != nil {
		return err
	}

	if DEBUG {
		fmt.Printf("txLink parent: %d child: %d\n", parentID, childID)
	}

	return nil

}

func (s *DB) Unlink(src DBRef, dest DBRef) error {

	return s.db.Update(func(tx *bolt.Tx) error {
		childBucket := tx.Bucket([]byte("child"))
		parentBucket := tx.Bucket([]byte("parent"))
		return s.txUnlink(childBucket, parentBucket, src, dest)
	})
}

func (s *DB) txUnlink(childBucket *bolt.Bucket, parentBucket *bolt.Bucket, parentID DBRef, childID DBRef) error {

	err := childBucket.Delete(parentID.ChildKey(childID))
	if err != nil {
		return err
	}

	err = parentBucket.Delete(childID.Key())
	if err != nil {
		return err
	}

	if DEBUG {
		fmt.Printf("txUnlink parent: %d child: %d\n", parentID, childID)
	}

	return nil

}

func (s *DB) txDelete(b *bolt.Bucket, id DBRef) error {
	if DEBUG {
		fmt.Printf("txDelete %d\n", id)
	}

	return b.Delete(id.Key())
}