From ea6a7b1d5e52a3550b4eb96e835b00f0e1841f33 Mon Sep 17 00:00:00 2001 From: Thomas Jungblut Date: Fri, 9 Sep 2022 12:26:43 +0200 Subject: [PATCH 1/2] add etcd --- README.md | 11 +++ cmd/go-ycsb/main.go | 6 +- db/etcd/db.go | 163 ++++++++++++++++++++++++++++++++++++++++++++ db/etcd/doc.go | 3 + 4 files changed, 181 insertions(+), 2 deletions(-) create mode 100644 db/etcd/db.go create mode 100644 db/etcd/doc.go diff --git a/README.md b/README.md index f64db12c..e7184a27 100644 --- a/README.md +++ b/README.md @@ -299,6 +299,17 @@ Common configurations: |bolt.mmap_flags|0|Set the DB.MmapFlags flag before memory mapping the file| |bolt.initial_mmap_size|0|The initial mmap size of the database in bytes. If <= 0, the initial map size is 0. If the size is smaller than the previous database, it takes no effect| +### etcd + +|field|default value|description| +|-|-|-| +|etcd.endpoints|"localhost:2379"|The etcd endpoint(s), multiple endpoints can be passed separated by comma.| +|etcd.dial_timeout|"2s"|The dial timeout duration passed into the client config.| +|etcd.cert_file|""|When using secure etcd, this should point to the crt file.| +|etcd.key_file|""|When using secure etcd, this should point to the pem file.| +|etcd.cacert_file|""|When using secure etcd, this should point to the ca file.| + + ## TODO - [ ] Support more measurement, like HdrHistogram diff --git a/cmd/go-ycsb/main.go b/cmd/go-ycsb/main.go index 28a8ceba..9eb52ddb 100644 --- a/cmd/go-ycsb/main.go +++ b/cmd/go-ycsb/main.go @@ -15,6 +15,7 @@ package main import ( "context" + "fmt" "log" "net/http" _ "net/http/pprof" @@ -28,7 +29,7 @@ import ( // Register workload - "fmt" + "github.com/spf13/cobra" "github.com/pingcap/go-ycsb/pkg/client" "github.com/pingcap/go-ycsb/pkg/measurement" @@ -36,7 +37,6 @@ import ( "github.com/pingcap/go-ycsb/pkg/util" _ "github.com/pingcap/go-ycsb/pkg/workload" "github.com/pingcap/go-ycsb/pkg/ycsb" - "github.com/spf13/cobra" // Register basic database _ "github.com/pingcap/go-ycsb/db/basic" @@ -70,6 +70,8 @@ import ( _ "github.com/pingcap/go-ycsb/db/minio" // Register elastic _ "github.com/pingcap/go-ycsb/db/elasticsearch" + // Register etcd + _ "github.com/pingcap/go-ycsb/db/etcd" ) var ( diff --git a/db/etcd/db.go b/db/etcd/db.go new file mode 100644 index 00000000..10580a84 --- /dev/null +++ b/db/etcd/db.go @@ -0,0 +1,163 @@ +package etcd + +import ( + "bytes" + "context" + "crypto/tls" + "encoding/json" + "fmt" + "strings" + "time" + + clientv3 "go.etcd.io/etcd/client/v3" + + "github.com/magiconair/properties" + "go.etcd.io/etcd/client/pkg/v3/transport" + + "github.com/pingcap/go-ycsb/pkg/ycsb" +) + +// properties +const ( + etcdEndpoints = "etcd.endpoints" + etcdDialTimeout = "etcd.dial_timeout" + etcdCertFile = "etcd.cert_file" + etcdKeyFile = "etcd.key_file" + etcdCaFile = "etcd.cacert_file" +) + +type etcdCreator struct{} + +type etcdDB struct { + p *properties.Properties + client *clientv3.Client +} + +func init() { + ycsb.RegisterDBCreator("etcd", etcdCreator{}) +} + +func (c etcdCreator) Create(p *properties.Properties) (ycsb.DB, error) { + cfg, err := getClientConfig(p) + if err != nil { + return nil, err + } + + client, err := clientv3.New(*cfg) + if err != nil { + return nil, err + } + + return &etcdDB{ + p: p, + client: client, + }, nil +} + +func getClientConfig(p *properties.Properties) (*clientv3.Config, error) { + endpoints := p.GetString(etcdEndpoints, "localhost:2379") + dialTimeout := p.GetDuration(etcdDialTimeout, 2*time.Second) + + var tlsConfig *tls.Config + if strings.Contains(endpoints, "https") { + tlsInfo := transport.TLSInfo{ + CertFile: p.MustGetString(etcdCertFile), + KeyFile: p.MustGetString(etcdKeyFile), + TrustedCAFile: p.MustGetString(etcdCaFile), + } + c, err := tlsInfo.ClientConfig() + if err != nil { + return nil, err + } + tlsConfig = c + } + + return &clientv3.Config{ + Endpoints: strings.Split(endpoints, ","), + DialTimeout: dialTimeout, + TLS: tlsConfig, + }, nil +} + +func (db *etcdDB) Close() error { + return db.client.Close() +} + +func (db *etcdDB) InitThread(ctx context.Context, _ int, _ int) context.Context { + return ctx +} + +func (db *etcdDB) CleanupThread(_ context.Context) { +} + +func getRowKey(table string, key string) string { + return fmt.Sprintf("%s:%s", table, key) +} + +func (db *etcdDB) Read(ctx context.Context, table string, key string, _ []string) (map[string][]byte, error) { + rkey := getRowKey(table, key) + value, err := db.client.Get(ctx, rkey) + if err != nil { + return nil, err + } + + if value.Count == 0 { + return nil, fmt.Errorf("could not find value for key [%s]", rkey) + } + + var r map[string][]byte + err = json.NewDecoder(bytes.NewReader(value.Kvs[0].Value)).Decode(&r) + if err != nil { + return nil, err + } + return r, nil +} + +func (db *etcdDB) Scan(ctx context.Context, table string, startKey string, count int, _ []string) ([]map[string][]byte, error) { + res := make([]map[string][]byte, count) + rkey := getRowKey(table, startKey) + values, err := db.client.Get(ctx, rkey, clientv3.WithFromKey(), clientv3.WithLimit(int64(count))) + if err != nil { + return nil, err + } + + if values.Count != int64(count) { + return nil, fmt.Errorf("unexpected number of result for key [%s], expected %d but was %d", rkey, count, values.Count) + } + + for _, v := range values.Kvs { + var r map[string][]byte + err = json.NewDecoder(bytes.NewReader(v.Value)).Decode(&r) + if err != nil { + return nil, err + } + res = append(res, r) + } + return res, nil +} + +func (db *etcdDB) Update(ctx context.Context, table string, key string, values map[string][]byte) error { + rkey := getRowKey(table, key) + data, err := json.Marshal(values) + if err != nil { + return err + } + _, err = db.client.Put(ctx, rkey, string(data)) + if err != nil { + return err + } + + return nil +} + +func (db *etcdDB) Insert(ctx context.Context, table string, key string, values map[string][]byte) error { + return db.Update(ctx, table, key, values) +} + +func (db *etcdDB) Delete(ctx context.Context, table string, key string) error { + _, err := db.client.Delete(ctx, getRowKey(table, key)) + if err != nil { + return err + } + return nil +} diff --git a/db/etcd/doc.go b/db/etcd/doc.go new file mode 100644 index 00000000..3ae594e7 --- /dev/null +++ b/db/etcd/doc.go @@ -0,0 +1,3 @@ +package etcd + +// If you want to use etcd, please follow the [Getting Started](https://github.com/etcd-io/etcd#getting-etcd) guide to install it. From 45f105c7fdf1759354aae13f71b51bb058e25bd9 Mon Sep 17 00:00:00 2001 From: Thomas Jungblut Date: Mon, 12 Sep 2022 17:00:31 +0200 Subject: [PATCH 2/2] readme --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index e7184a27..af87a49d 100644 --- a/README.md +++ b/README.md @@ -98,6 +98,7 @@ Available Commands: - MongoDB - Redis and Redis Cluster - BoltDB +- etcd ## Database Configuration