Switch to BoltDB

bind.go

@@ -5,7 +5,7 @@ import (
 	"strings"
 
 	"github.com/bwmarrin/discordgo"
-	"github.com/syndtr/goleveldb/leveldb"
+	"go.etcd.io/bbolt"
 )
 
 func bind(player string, user *discordgo.User) (id uint32, err error) {
@@ -18,43 +18,34 @@ func bind(player string, user *discordgo.User) (id uint32, err error) {
 }
 
 func putBind(playerID uint32, discordUser *discordgo.User) (err error) {
-	var tx *leveldb.Transaction
-	tx, err = db.OpenTransaction()
-	if err != nil {
-		return
-	}
-	defer commitOrDiscard(tx, &err)
-	if err = putUInt32(tx, normalPath, discordUser.String(), playerID); err != nil {
-		return
-	}
-	if err = putLowercaseIndex(tx, discordUser.String()); err != nil {
-		return
-	}
-	return
-}
-
-func putLowercaseIndex(tx *leveldb.Transaction, username string) error {
-	return putString(tx, lowercasePath, strings.ToLower(username), username)
+	return bdb.Update(func(t *bbolt.Tx) (err error) {
+		name := discordUser.String()
+		err = newUsersBucket(t).put(name, playerID)
+		if err != nil {
+			return
+		}
+		return newLowercaseBucket(t).put(name)
+	})
 }
 
 func getBind(username string) (playerID uint32, err error) {
-	playerID, err = getUInt32(normalPath, username)
-	if err == leveldb.ErrNotFound {
-		return 0, fmt.Errorf("player %s isn't in the database. Use `-bind <Steam ID>` to register", username)
-	}
-	if err != nil {
+	err = bdb.View(func(t *bbolt.Tx) (err error) {
+		playerID, err = newUsersBucket(t).get(username)
 		return
+	})
+	if err != nil {
+		return 0, fmt.Errorf("player %s isn't in the database. Use `-bind <Steam ID>` to register", username)
 	}
 	return
 }
 
 func deleteBind(user *discordgo.User) (err error) {
-	var tx *leveldb.Transaction
-	tx, err = db.OpenTransaction()
-	if err != nil {
+	return bdb.Update(func(t *bbolt.Tx) (err error) {
+		err = newUsersBucket(t).del(user.String())
+		if err != nil {
+			return
+		}
+		err = newLowercaseBucket(t).del(strings.ToLower(user.String()))
 		return
-	}
-	defer commitOrDiscard(tx, &err)
-	err = deleteString(tx, normalPath, user.String())
-	return
+	})
 }

boltdb.go

@@ -0,0 +1,91 @@
+package main
+
+import (
+	"bytes"
+	"fmt"
+	"strings"
+
+	"go.etcd.io/bbolt"
+)
+
+var (
+	bdb                 *bbolt.DB
+	discordBucketName   = []byte("discord_to_steamid")
+	lowercaseBucketName = []byte("lowercase_to_normalcase")
+)
+
+type usersBucket struct {
+	*bbolt.Bucket
+}
+
+func newUsersBucket(tx *bbolt.Tx) usersBucket {
+	return usersBucket{tx.Bucket(discordBucketName)}
+}
+
+func (b usersBucket) put(key string, value uint32) error {
+	return b.Put([]byte(key), uint32ToBytes(value))
+}
+
+func (b usersBucket) get(key string) (uint32, error) {
+	val := b.Get([]byte(key))
+	if val == nil {
+		return 0, fmt.Errorf("not found")
+	}
+	return uint32FromBytes(val), nil
+}
+
+func (b usersBucket) del(key string) error {
+	return b.Delete([]byte(key))
+}
+
+type lowercaseBucket struct {
+	*bbolt.Bucket
+}
+
+func newLowercaseBucket(tx *bbolt.Tx) lowercaseBucket {
+	return lowercaseBucket{tx.Bucket(lowercaseBucketName)}
+}
+
+func (l lowercaseBucket) put(name string) error {
+	return l.Put([]byte(strings.ToLower(name)), []byte(name))
+}
+
+func (l lowercaseBucket) get(key string) string {
+	return string(l.Get([]byte(key)))
+}
+
+func (l lowercaseBucket) findFirstString(prefix string) (result string, err error) {
+	c := l.Cursor()
+	k, v := c.Seek([]byte(prefix))
+	if k == nil || !bytes.HasPrefix(k, []byte(prefix)) {
+		return "", fmt.Errorf("not found")
+	}
+	return string(v), nil
+}
+
+func (l lowercaseBucket) del(key string) error {
+	return l.Delete([]byte(key))
+}
+
+func openBoltDB(dbPath string) (err error) {
+	bdb, err = bbolt.Open(dbPath, 0600, nil)
+	return
+}
+
+func initBoltDB() {
+	err := bdb.Update(func(t *bbolt.Tx) error {
+		_, err := t.CreateBucketIfNotExists(discordBucketName)
+		if err != nil {
+			return err
+		}
+		_, err = t.CreateBucketIfNotExists(lowercaseBucketName)
+		return err
+	})
+	if err != nil {
+		panic(err)
+	}
+}
+
+func closeBoltDB() error {
+	return bdb.Close()
+}

bot.go

@@ -17,6 +17,7 @@ import (
 
 	"github.com/Philipp15b/go-steamapi"
 	"github.com/bwmarrin/discordgo"
+	"go.etcd.io/bbolt"
 )
 
 const (
@@ -41,7 +42,11 @@ var (
 )
 
 func playerIDFromDiscordName(username string) (uint32, error) {
-	discordName, err := findFirstString(lowercasePath, username)
+	var discordName string
+	err := bdb.View(func(t *bbolt.Tx) (err error) {
+		discordName, err = newLowercaseBucket(t).findFirstString(username)
+		return
+	})
 	if err != nil {
 		return 0, fmt.Errorf("discord user name starting with '%s' was not found", username)
 	}
@@ -219,7 +224,7 @@ func bot() (err error) {
 	select {
 	case <-sc:
 	case <-restartChan:
-		db.Close()
+		ldb.Close()
 		cmd := exec.Command(os.Args[0], os.Args[1:]...)
 		log.Println("Restarting myself...")
 		err := cmd.Start()

config.go

@@ -69,7 +69,8 @@ var config struct {
 	Token         string
 	SteamKey      string        `json:"steam_key"`
 	ChannelID     string        `json:"channel_id"`
-	DBPath        string        `json:"database_path"`
+	LevelDBPath   string        `json:"ldb_database_path"`
+	BoltDBPath    string        `json:"bdb_database_path"`
 	QueryInterval time.Duration `json:"query_interval"`
 	Servers       []*ns2server
 	Seeding       seeding
@@ -95,7 +96,7 @@ func loadConfigFilename(filename string) error {
 	if config.Token == "" {
 		return fmt.Errorf("specify token in config.json")
 	}
-	if config.DBPath == "" {
+	if config.LevelDBPath == "" {
 		return fmt.Errorf("specify database_path in config.json")
 	}
 	return nil

go.mod

@@ -8,6 +8,7 @@ require (
 	github.com/docopt/docopt-go v0.0.0-20180111231733-ee0de3bc6815
 	github.com/rumblefrog/go-a2s v1.0.0
 	github.com/syndtr/goleveldb v1.0.0
+	go.etcd.io/bbolt v1.3.5
 )
 
 // temporary override until pull requests #11 and #12 are merged

go.sum

@@ -23,6 +23,8 @@ github.com/rumblefrog/go-a2s v1.0.0 h1:9IIVIOQ1bXZJeTilmzkJDeGa/9W1c089VciTbp+Wp
 github.com/rumblefrog/go-a2s v1.0.0/go.mod h1:JwbTgMTRGZcWzr3T2MUfDusrJU5Bdg8biEeZzPtN0So=
 github.com/syndtr/goleveldb v1.0.0 h1:fBdIW9lB4Iz0n9khmH8w27SJ3QEJ7+IgjPEwGSZiFdE=
 github.com/syndtr/goleveldb v1.0.0/go.mod h1:ZVVdQEZoIme9iO1Ch2Jdy24qqXrMMOU6lpPAyBWyWuQ=
+go.etcd.io/bbolt v1.3.5 h1:XAzx9gjCb0Rxj7EoqcClPD1d5ZBxZJk0jbuoPHenBt0=
+go.etcd.io/bbolt v1.3.5/go.mod h1:G5EMThwa9y8QZGBClrRx5EY+Yw9kAhnjy3bSjsnlVTQ=
 golang.org/x/crypto v0.0.0-20181030102418-4d3f4d9ffa16 h1:y6ce7gCWtnH+m3dCjzQ1PCuwl28DDIc3VNnvY29DlIA=
 golang.org/x/crypto v0.0.0-20181030102418-4d3f4d9ffa16/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4=
 golang.org/x/net v0.0.0-20180906233101-161cd47e91fd h1:nTDtHvHSdCn1m6ITfMRqtOd/9+7a3s8RBNOZ3eYZzJA=
@@ -31,6 +33,7 @@ golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f h1:wMNYb4v58l5UBM7MYRLPG6Zh
 golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e h1:o3PsSEY8E4eXWkXrIP9YJALUkVZqzHJT5DOasTyn8Vs=
 golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
+golang.org/x/sys v0.0.0-20200202164722-d101bd2416d5/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/text v0.3.0 h1:g61tztE5qeGQ89tm6NTjjM9VPIm088od1l6aSorWRWg=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=

db.go → leveldb.go

@@ -12,7 +12,7 @@ import (
 
 const pathSeparator = "\x00"
 
-var db *leveldb.DB
+var ldb *leveldb.DB
 
 // commitOrDiscard either commits or discard the transaction depending on the error presence.
 // This function is supposed to be called in defer.
@@ -27,16 +27,16 @@ func commitOrDiscard(tx *leveldb.Transaction, err *error) {
 }
 
 func openDB(dbPath string) (err error) {
-	db, err = leveldb.OpenFile(dbPath, nil)
+	ldb, err = leveldb.OpenFile(dbPath, nil)
 	if e, corrupted := err.(*errors.ErrCorrupted); corrupted {
 		log.Printf("WARNING: database corruption: %s. Attempting to recover...", e)
-		db, err = leveldb.RecoverFile(dbPath, nil)
+		ldb, err = leveldb.RecoverFile(dbPath, nil)
 	}
 	return
 }
 
 func closeDB() error {
-	return db.Close()
+	return ldb.Close()
 }
 
 func makePath(path ...string) string {
@@ -51,7 +51,7 @@ func pathKey(path string, key string) []byte {
 }
 
 func findFirstString(path string, prefix string) (result string, err error) {
-	iter := db.NewIterator(util.BytesPrefix(pathKey(path, prefix)), nil)
+	iter := ldb.NewIterator(util.BytesPrefix(pathKey(path, prefix)), nil)
 	defer iter.Release()
 	if ok := iter.Next(); ok {
 		return string(iter.Value()), nil
@@ -60,7 +60,7 @@ func findFirstString(path string, prefix string) (result string, err error) {
 }
 
 func getUInt32(path string, key string) (uint32, error) {
-	val, err := db.Get(pathKey(path, key), nil)
+	val, err := ldb.Get(pathKey(path, key), nil)
 	if err != nil {
 		return 0, err
 	}
@@ -88,3 +88,7 @@ func uint32ToBytes(val uint32) []byte {
 func putUInt32(tx *leveldb.Transaction, path string, key string, value uint32) error {
 	return tx.Put(pathKey(path, key), uint32ToBytes(value), nil)
 }
+
+func putLowercaseIndex(tx *leveldb.Transaction, username string) error {
+	return putString(tx, lowercasePath, strings.ToLower(username), username)
+}

main.go

@@ -11,11 +11,13 @@ func main() {
 	ns2query [-c config]
 	ns2query -u
 	ns2query -h
+	ns2query --convert
 
 Options:
 	-h --help    This help
 	-c config    Use config file [default: config.json]
 	-u           Update database
+	--convert    Convert the database to BoltDB
 `
 	opts, err := docopt.ParseDoc(usage)
 	if err != nil {
@@ -24,14 +26,31 @@ Options:
 	if err := loadConfig(opts["-c"].(string)); err != nil {
 		log.Fatal("error loading config:", err)
 	}
-	if err := openDB(config.DBPath); err != nil {
-		log.Fatal("error opening database:", err)
+	if config.LevelDBPath != "" {
+		if err := openDB(config.LevelDBPath); err != nil {
+			log.Fatal("error opening LevelDB database:", err)
+		}
+		defer closeDB()
+	}
+	if config.BoltDBPath != "" {
+		if err := openBoltDB(config.BoltDBPath); err != nil {
+			log.Fatal("error opening BoltDB database:", err)
+		}
+		initBoltDB()
+		defer closeBoltDB()
 	}
-	defer closeDB()
 	if update, err := opts.Bool("-u"); err == nil && update {
 		if err := updateDB(); err != nil {
 			log.Fatal("error updating db:", err)
 		}
+		log.Println("Database updated successfully.")
+		return
+	}
+	if convert, _ := opts.Bool("--convert"); convert {
+		if err := convertDB(); err != nil {
+			log.Fatal("error converting db:", err)
+		}
+		log.Println("Database converted successfully.")
 		return
 	}
 	err = bot()

updatedb.go

@@ -4,11 +4,13 @@ import (
 	"strings"
 
 	"github.com/syndtr/goleveldb/leveldb"
+	"github.com/syndtr/goleveldb/leveldb/util"
+	"go.etcd.io/bbolt"
 )
 
 func updateDB() (err error) {
 	var tx *leveldb.Transaction
-	tx, err = db.OpenTransaction()
+	tx, err = ldb.OpenTransaction()
 	if err != nil {
 		return
 	}
@@ -31,3 +33,20 @@ func updateDB() (err error) {
 	}
 	return
 }
+
+func convertDB() (err error) {
+	err = bdb.Update(func(t *bbolt.Tx) (err error) {
+		users := newUsersBucket(t)
+		lc := newLowercaseBucket(t)
+		path := normalPath + "\x00"
+		iter := ldb.NewIterator(util.BytesPrefix([]byte(path)), nil)
+		defer iter.Release()
+		for iter.Next() {
+			name := strings.TrimPrefix(string(iter.Key()), path)
+			users.put(name, uint32FromBytes(iter.Value()))
+			lc.put(name)
+		}
+		return
+	})
+	return
+}