2019-02-13 22:45:23 +08:00
|
|
|
package main
|
|
|
|
|
|
|
|
import (
|
|
|
|
"encoding/json"
|
|
|
|
"log"
|
|
|
|
"net/http"
|
2019-02-18 23:43:49 +08:00
|
|
|
"os"
|
2019-02-13 22:45:23 +08:00
|
|
|
|
2019-06-23 08:44:23 +08:00
|
|
|
. "store/backend-protobuf/go"
|
|
|
|
|
2019-02-18 23:43:49 +08:00
|
|
|
"github.com/joho/godotenv"
|
2019-02-13 22:45:23 +08:00
|
|
|
"github.com/dgraph-io/badger"
|
|
|
|
"github.com/nats-io/go-nats"
|
|
|
|
"github.com/golang/protobuf/proto"
|
|
|
|
)
|
|
|
|
|
|
|
|
var dbPath string
|
|
|
|
var natsHost string
|
|
|
|
|
|
|
|
var db *badger.DB
|
|
|
|
var nc *nats.Conn
|
|
|
|
|
|
|
|
func main() {
|
2019-02-18 23:43:49 +08:00
|
|
|
// Load .env
|
|
|
|
err := godotenv.Load()
|
|
|
|
if err != nil {
|
|
|
|
log.Fatal("Error loading .env file")
|
|
|
|
}
|
|
|
|
dbPath = os.Getenv("DBPATH")
|
|
|
|
natsHost = os.Getenv("NATS")
|
2019-02-13 22:45:23 +08:00
|
|
|
|
|
|
|
// Open badger
|
|
|
|
log.Printf("starting badger at %s", dbPath)
|
|
|
|
opts := badger.DefaultOptions
|
|
|
|
opts.Dir = dbPath
|
|
|
|
opts.ValueDir = dbPath
|
|
|
|
db, err = badger.Open(opts)
|
|
|
|
if err != nil {
|
|
|
|
log.Fatal(err)
|
|
|
|
}
|
|
|
|
defer db.Close()
|
|
|
|
|
|
|
|
// NATS client
|
|
|
|
nc, err = nats.Connect(natsHost)
|
|
|
|
if err != nil {
|
|
|
|
log.Fatal(err)
|
|
|
|
}
|
|
|
|
nc.Subscribe("new_store", NewStore)
|
|
|
|
|
|
|
|
nc.Subscribe("request_store", RequestStore)
|
|
|
|
nc.Subscribe("scan_store", ScanStore)
|
|
|
|
defer nc.Close()
|
|
|
|
|
|
|
|
select { } // Wait forever
|
|
|
|
}
|
|
|
|
|
|
|
|
func NewStore(m *nats.Msg) {
|
|
|
|
storeRequest := Store{}
|
|
|
|
if err := proto.Unmarshal(m.Data, &storeRequest); err != nil {
|
2019-02-16 01:29:29 +08:00
|
|
|
log.Println(err) // Fail quietly since protobuf data is needed torespond
|
2019-02-13 22:45:23 +08:00
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
key, err := MarshalKey(storeRequest.Type, storeRequest.Bite.Key, storeRequest.Bite.Start)
|
|
|
|
if err != nil {
|
|
|
|
log.Println(err)
|
2019-02-16 01:29:29 +08:00
|
|
|
errRes := Response {
|
|
|
|
Code: 400,
|
2019-02-18 03:10:00 +08:00
|
|
|
Message: []byte(http.StatusText(http.StatusBadRequest)),
|
2019-02-16 01:29:29 +08:00
|
|
|
Client: storeRequest.Bite.Client,
|
|
|
|
}
|
|
|
|
errResBytes, errResErr := proto.Marshal(&errRes)
|
|
|
|
if errResErr == nil {
|
|
|
|
nc.Publish("res", errResBytes)
|
|
|
|
}
|
2019-02-13 22:45:23 +08:00
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
err = db.Update(func(txn *badger.Txn) error {
|
|
|
|
// TODO: prevent overwriting existing
|
|
|
|
err := txn.Set(key, storeRequest.Bite.Data)
|
|
|
|
return err
|
|
|
|
})
|
|
|
|
|
|
|
|
if err != nil {
|
|
|
|
log.Println(err)
|
2019-02-16 01:29:29 +08:00
|
|
|
errRes := Response {
|
|
|
|
Code: 500,
|
2019-02-18 03:10:00 +08:00
|
|
|
Message: []byte(http.StatusText(http.StatusInternalServerError)),
|
2019-02-16 01:29:29 +08:00
|
|
|
Client: storeRequest.Bite.Client,
|
|
|
|
}
|
|
|
|
errResBytes, errResErr := proto.Marshal(&errRes)
|
|
|
|
if errResErr == nil {
|
|
|
|
nc.Publish("res", errResBytes)
|
|
|
|
}
|
2019-02-13 22:45:23 +08:00
|
|
|
return
|
2019-02-16 01:29:29 +08:00
|
|
|
} else {
|
|
|
|
res := Response {
|
|
|
|
Code: 200,
|
|
|
|
Message: []byte(key),
|
|
|
|
Client: storeRequest.Bite.Client,
|
|
|
|
}
|
|
|
|
resBytes, resErr := proto.Marshal(&res)
|
|
|
|
if resErr == nil {
|
|
|
|
nc.Publish("res", resBytes)
|
|
|
|
}
|
2019-02-13 22:45:23 +08:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func RequestStore(m *nats.Msg) {
|
|
|
|
req := DataRequest{}
|
|
|
|
if err := proto.Unmarshal(m.Data, &req); err != nil {
|
|
|
|
log.Println(err)
|
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
key, err := MarshalKey(req.Type, req.Key, req.Start)
|
|
|
|
|
|
|
|
err = db.View(func(txn *badger.Txn) error {
|
|
|
|
item, err := txn.Get(key)
|
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
2019-02-18 03:10:00 +08:00
|
|
|
value, err := item.Value()
|
2019-02-13 22:45:23 +08:00
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
2019-02-18 03:10:00 +08:00
|
|
|
|
|
|
|
res := Response {
|
|
|
|
Code: 200,
|
|
|
|
Message: value,
|
|
|
|
}
|
|
|
|
resBytes, err := proto.Marshal(&res)
|
|
|
|
|
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
|
|
|
|
nc.Publish(m.Reply, resBytes)
|
2019-02-13 22:45:23 +08:00
|
|
|
return nil
|
|
|
|
})
|
|
|
|
|
|
|
|
if err != nil {
|
|
|
|
res := ReplyError(err.Error(), 400)
|
|
|
|
nc.Publish(m.Reply, res)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
type BitesList struct {
|
|
|
|
Previous uint64 `json:"previous"` // One bite before starts. Hint for how many steps the client can skip
|
|
|
|
Starts []uint64 `json:"starts"`
|
|
|
|
Next uint64 `json:"next"` // One bite after starts. Hint for how many steps the client can skip
|
|
|
|
}
|
|
|
|
|
|
|
|
func ScanStore(m *nats.Msg) {
|
|
|
|
req := ScanRequest {}
|
|
|
|
if err := proto.Unmarshal(m.Data, &req); err != nil {
|
|
|
|
log.Println(err)
|
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
prefix, err := MarshalKeyPrefix(req.Type, req.Key)
|
|
|
|
|
|
|
|
if err != nil {
|
|
|
|
res := ReplyError(http.StatusText(http.StatusBadRequest), http.StatusBadRequest)
|
|
|
|
nc.Publish(m.Reply, res)
|
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
fromKey, err := MarshalKey(req.Type, req.Key, req.From)
|
|
|
|
if err != nil {
|
|
|
|
res := ReplyError(http.StatusText(http.StatusBadRequest), http.StatusBadRequest)
|
|
|
|
nc.Publish(m.Reply, res)
|
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
bitesList := BitesList {}
|
|
|
|
|
|
|
|
err = db.View(func(txn *badger.Txn) error {
|
|
|
|
opts := badger.DefaultIteratorOptions
|
|
|
|
opts.PrefetchValues = false
|
|
|
|
opts.Reverse = true
|
|
|
|
it := txn.NewIterator(opts)
|
|
|
|
defer it.Close()
|
|
|
|
|
|
|
|
// Fetch previous key
|
|
|
|
it.Seek(fromKey)
|
|
|
|
if it.ValidForPrefix(fromKey) {
|
|
|
|
// Lazy check to compare key == seeked key
|
|
|
|
it.Next()
|
|
|
|
}
|
|
|
|
if !it.ValidForPrefix(prefix) {
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
item := it.Item()
|
|
|
|
key := item.Key()
|
|
|
|
|
|
|
|
_, _, start, err := ExtractKey(key)
|
|
|
|
if err != nil {
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
bitesList.Previous = start
|
|
|
|
|
|
|
|
return nil
|
|
|
|
})
|
|
|
|
if err != nil {
|
|
|
|
res := ReplyError(err.Error(), http.StatusBadRequest)
|
|
|
|
nc.Publish(m.Reply, res)
|
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
err = db.View(func(txn *badger.Txn) error {
|
|
|
|
opts := badger.DefaultIteratorOptions
|
|
|
|
opts.PrefetchValues = false
|
|
|
|
it := txn.NewIterator(opts)
|
|
|
|
defer it.Close()
|
|
|
|
|
|
|
|
for it.Seek(fromKey); it.ValidForPrefix(prefix); it.Next() {
|
|
|
|
item := it.Item()
|
|
|
|
key := item.Key()
|
|
|
|
|
|
|
|
_, _, start, err := ExtractKey(key)
|
|
|
|
if err != nil {
|
|
|
|
continue
|
|
|
|
}
|
|
|
|
if start > req.To {
|
|
|
|
// A key was found that is greater than to
|
|
|
|
// Save that as next
|
|
|
|
bitesList.Next = start
|
|
|
|
break
|
|
|
|
}
|
|
|
|
|
|
|
|
bitesList.Starts = append(bitesList.Starts, start)
|
|
|
|
}
|
|
|
|
|
|
|
|
return nil
|
|
|
|
})
|
|
|
|
if err != nil {
|
|
|
|
res := ReplyError(err.Error(), http.StatusBadRequest)
|
|
|
|
nc.Publish(m.Reply, res)
|
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
jsonString, err := json.Marshal(&bitesList)
|
|
|
|
res := Response {
|
|
|
|
Code: 200,
|
|
|
|
Message: []byte(jsonString),
|
|
|
|
}
|
|
|
|
resBytes, _ := proto.Marshal(&res)
|
|
|
|
nc.Publish(m.Reply, resBytes)
|
|
|
|
}
|
|
|
|
|
|
|
|
func ReplyError(msg string, code uint32) []byte {
|
|
|
|
res := Response {
|
|
|
|
Code: code,
|
|
|
|
Message: []byte(msg),
|
|
|
|
}
|
|
|
|
resBytes, _ := proto.Marshal(&res)
|
|
|
|
|
|
|
|
return resBytes
|
|
|
|
}
|