127 lines
2.3 KiB
Go
127 lines
2.3 KiB
Go
package db
|
|
|
|
import (
|
|
"errors"
|
|
"sync"
|
|
"time"
|
|
|
|
schemav1 "git.giftfish.de/ston1th/haproxy-lb/pkg/api/v1/schema"
|
|
"git.giftfish.de/ston1th/haproxy-lb/pkg/cluster"
|
|
)
|
|
|
|
type DB struct {
|
|
sync.Mutex
|
|
kv cluster.KV
|
|
}
|
|
|
|
func New(kv cluster.KV) *DB {
|
|
return &DB{kv: kv}
|
|
}
|
|
|
|
const (
|
|
ipPrefix = "/ip/"
|
|
delPrefix = "/del/"
|
|
namePrefix = "/name/"
|
|
lbPrefix = "/lb/"
|
|
)
|
|
|
|
func (db *DB) GetIPs() (ips []string, err error) {
|
|
m, err := db.kv.GetPrefix(ipPrefix)
|
|
if err != nil {
|
|
return
|
|
}
|
|
for k := range m {
|
|
ips = append(ips, k)
|
|
}
|
|
return
|
|
}
|
|
func (db *DB) GetDeletedIPs() (ips []string, err error) {
|
|
err = timeout(func() error {
|
|
m, err := db.kv.GetPrefix(delPrefix)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
for k := range m {
|
|
ips = append(ips, k)
|
|
db.kv.Delete(delPrefix + k)
|
|
}
|
|
return nil
|
|
})
|
|
return
|
|
}
|
|
|
|
func timeout(f func() error) error {
|
|
return timeoutWithDuration(f, time.Second*3)
|
|
}
|
|
func timeoutWithDuration(f func() error, d time.Duration) error {
|
|
c := make(chan error, 1)
|
|
go func() {
|
|
e := f()
|
|
c <- e
|
|
}()
|
|
select {
|
|
case err := <-c:
|
|
return err
|
|
case <-time.After(d):
|
|
return errors.New("timeout")
|
|
}
|
|
}
|
|
|
|
func (db *DB) GetCIDR(name string) (cidr string, err error) {
|
|
b, err := db.GetLB(name)
|
|
if err != nil {
|
|
return
|
|
}
|
|
lb, err := schemav1.NewLoadBalancerFromBytes(b)
|
|
if err != nil {
|
|
return
|
|
}
|
|
cidr = lb.CIDR
|
|
return
|
|
}
|
|
|
|
func (db *DB) GetLBs(cl string) (lbs map[string][]byte, err error) {
|
|
lbs = make(map[string][]byte)
|
|
if cl != "" {
|
|
cl = cl + "/"
|
|
}
|
|
m, err := db.kv.GetPrefix(lbPrefix + cl)
|
|
if err != nil {
|
|
return
|
|
}
|
|
for k, v := range m {
|
|
lbs[k] = v
|
|
}
|
|
return
|
|
}
|
|
func (db *DB) GetLB(name string) ([]byte, error) {
|
|
return db.kv.Get(lbPrefix + name)
|
|
}
|
|
func (db *DB) LBExists(name string) bool {
|
|
_, err := db.GetLB(name)
|
|
return err == nil
|
|
}
|
|
func (db *DB) SetLB(name string, lb []byte) error {
|
|
return db.kv.Set(lbPrefix+name, lb)
|
|
}
|
|
func (db *DB) DeleteLB(name string) error {
|
|
return db.kv.Delete(lbPrefix + name)
|
|
}
|
|
|
|
func (db *DB) Exists(key string) (bool, error) {
|
|
_, err := db.kv.Get(ipPrefix + key)
|
|
if err == cluster.ErrKeyNotFound {
|
|
return false, nil
|
|
}
|
|
return err == nil, err
|
|
}
|
|
func (db *DB) Add(key string) error {
|
|
return db.kv.Set(ipPrefix+key, nil)
|
|
}
|
|
func (db *DB) Delete(key string) error {
|
|
err := db.kv.Delete(ipPrefix + key)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return db.kv.Set(delPrefix+key, nil)
|
|
}
|