diff --git a/.gitignore b/.gitignore index 6b7c1cdd5..6d5b0c8fe 100644 --- a/.gitignore +++ b/.gitignore @@ -75,3 +75,11 @@ tags.dot # SQL3 /sql3/sql3.html + +staticcheck.conf + + +.quick +dax/dax-data + +coverage-from-docker \ No newline at end of file diff --git a/Dockerfile b/Dockerfile index 59f906052..0ca6ef677 100644 --- a/Dockerfile +++ b/Dockerfile @@ -38,7 +38,7 @@ FROM alpine:3.13.2 as runner LABEL maintainer "dev@molecula.com" -RUN apk add --no-cache curl jq +RUN apk add --no-cache curl jq tree COPY --from=pilosa-builder /pilosa/build/featurebase / diff --git a/Dockerfile-clustertests b/Dockerfile-clustertests index 17fbb891e..11160bbef 100644 --- a/Dockerfile-clustertests +++ b/Dockerfile-clustertests @@ -20,16 +20,19 @@ RUN apt install -y docker.io ADD https://github.com/docker/compose/releases/latest/download/docker-compose-Linux-x86_64 /usr/local/bin/docker-compose RUN chmod +x /usr/local/bin/docker-compose +WORKDIR /go/src/github.com/molecula/featurebase/cmd/featurebase + # generate an instrumented binary to allow for calculating code coverage for clustertests # the entrypoint for the binary is TestRunMain, which is wrapper for main -RUN cd /go/src/github.com/featurebasedb/featurebase/cmd/featurebase && \ - go test -covermode=atomic -coverpkg=../../... -c -tags testrunmain -o featurebase && \ - cp /go/src/github.com/featurebasedb/featurebase/cmd/featurebase/featurebase /featurebase +RUN go test -covermode=atomic -coverpkg=../../... -c -tags testrunmain -o featurebase +RUN cp /go/src/github.com/molecula/featurebase/cmd/featurebase/featurebase /featurebase COPY NOTICE /NOTICE EXPOSE 10101 VOLUME /data -ENTRYPOINT ["bash", "-c"] -CMD ["/featurebase", "-test.run=TestRunMain", "-test.coverprofile=/results/coverage.out", "server", "--data-dir", "/data", "--bind", "http://0.0.0.0:10101"] + +# use e.g. "-test.coverprofile=/results/coverage.out" +CMD ["/featurebase", "-test.run=TestRunMain", "server"] + diff --git a/Dockerfile-clustertests-client b/Dockerfile-clustertests-client index 1f3df34e9..b3290b032 100644 --- a/Dockerfile-clustertests-client +++ b/Dockerfile-clustertests-client @@ -19,9 +19,10 @@ RUN apt install -y docker.io ADD https://github.com/docker/compose/releases/latest/download/docker-compose-Linux-x86_64 /usr/local/bin/docker-compose RUN chmod +x /usr/local/bin/docker-compose -RUN cd /go/src/github.com/featurebasedb/featurebase/cmd/featurebase && \ - go test -covermode=atomic -coverpkg=../../... -c -tags testrunmain -o featurebase && \ - cp /go/src/github.com/featurebasedb/featurebase/cmd/featurebase/featurebase /featurebase +WORKDIR /go/src/github.com/molecula/featurebase/cmd/featurebase + +RUN go test -covermode=atomic -coverpkg=../../... -c -tags testrunmain -o featurebase +RUN cp /go/src/github.com/molecula/featurebase/cmd/featurebase/featurebase /featurebase COPY NOTICE /NOTICE @@ -31,5 +32,6 @@ COPY ./internal/clustertests /go/src/github.com/featurebasedb/featurebase/intern EXPOSE 10101 VOLUME /data -ENTRYPOINT ["bash", "-c"] +WORKDIR /go/src/github.com/molecula/featurebase + CMD ["/featurebase", "-test.run=TestRunMain", "-test.coverprofile=/results/coverage.out", "server", "--data-dir", "/data", "--bind", "http://0.0.0.0:10101"] diff --git a/Dockerfile-datagen b/Dockerfile-datagen new file mode 100644 index 000000000..6ae776463 --- /dev/null +++ b/Dockerfile-datagen @@ -0,0 +1,47 @@ +# syntax=docker/dockerfile:1 + +########################## +### datagen builder ### +########################## + +FROM golang:alpine as builder + +WORKDIR /featurebase + +COPY . ./ + +RUN apk add --no-cache build-base bash git make librdkafka pkgconfig + +# install librdkafka +RUN git clone https://github.com/edenhill/librdkafka.git +RUN cd librdkafka && ./configure --prefix /usr && make && make install + +ENV PKG_CONFIG_PATH=/usr/lib/pkgconfig/ + +RUN cd idk && make build-datagen + +# ENTRYPOINT ["tail", "-f", "/dev/null"] + +######################### +### datagen runner ### +######################### + +FROM alpine:3.15.3 as runner + +WORKDIR / + +LABEL maintainer "dev@molecula.com" + +RUN apk add --no-cache curl jq + +COPY --from=builder /featurebase/idk/build/datagen /bin/ +COPY --from=builder /usr/lib/librdkafka* /usr/lib/ +COPY idk/datagen/testdata/* /testdata/ + +EXPOSE 8080 + +# VOLUME /data +# ENV ADDR 0.0.0.0:8080 + +#ENTRYPOINT ["sleep", "infinity"] +ENTRYPOINT ["datagen"] diff --git a/Dockerfile-dax b/Dockerfile-dax new file mode 100644 index 000000000..6b5fbb9ed --- /dev/null +++ b/Dockerfile-dax @@ -0,0 +1,32 @@ +ARG GO_VERSION=latest + + +########################### +### FeatureBase Builder ### +########################### + +FROM golang:${GO_VERSION} as featurebase-builder +ARG MAKE_FLAGS +WORKDIR /fb + +COPY . ./ +RUN make build FLAGS="-o build/featurebase" ${MAKE_FLAGS} + +########################## +### FeatureBase runner ### +########################## + +FROM alpine:3.13.2 as runner + +LABEL maintainer "dev@featurebase.com" + +RUN apk add --no-cache curl jq tree + +COPY --from=featurebase-builder /fb/build/featurebase / + +COPY NOTICE /NOTICE + +EXPOSE 8080 + +ENTRYPOINT ["/featurebase"] +CMD ["dax"] diff --git a/Dockerfile-dax-quick b/Dockerfile-dax-quick new file mode 100644 index 000000000..e3e21a12f --- /dev/null +++ b/Dockerfile-dax-quick @@ -0,0 +1,18 @@ +ARG GO_VERSION=latest + +########################## +### FeatureBase runner ### +########################## + +FROM alpine:3.13.2 as runner + +LABEL maintainer "dev@featurebase.com" + +RUN apk add --no-cache curl jq tree + +COPY ./fb_linux /featurebase + +EXPOSE 8080 + +ENTRYPOINT ["/featurebase"] +CMD ["dax"] diff --git a/Makefile b/Makefile index 5d798782f..a97dad399 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,5 @@ -.PHONY: build clean build-lattice cover cover-viz default docker docker-build docker-tag-push generate generate-protoc generate-pql generate-statik generate-stringer install install-protoc-gen-gofast install-protoc install-statik install-peg test +.PHONY: build clean build-lattice cover cover-viz default docker docker-build docker-tag-push generate generate-protoc generate-pql generate-statik generate-stringer install install-protoc-gen-gofast install-protoc install-statik install-peg test docker-login + VERSION := $(shell git describe --tags 2> /dev/null || echo unknown) VARIANT = Molecula GO=go @@ -18,6 +19,7 @@ RACE_TEST_TIMEOUT=10m export GO111MODULE=on export GOPRIVATE=github.com/molecula export CGO_ENABLED=0 +AWS_ACCOUNTID ?= undefined # Run tests and compile Pilosa default: test build @@ -26,7 +28,7 @@ default: test build clean: rm -rf vendor build rm -f *.rpm *.deb - + # Set up vendor directory using `go mod vendor` vendor: go.mod $(GO) mod vendor @@ -92,10 +94,10 @@ build: $(GO) build -tags='$(BUILD_TAGS)' -ldflags $(LDFLAGS) $(FLAGS) ./cmd/featurebase package: - go build -o featurebase ./cmd/featurebase + GOOS=$(GOOS) GOARCH=$(GOARCH) FLAGS="-o featurebase" $(MAKE) build GOARCH=$(GOARCH) VERSION=$(VERSION) nfpm package --packager deb --target featurebase.$(VERSION).$(GOARCH).deb GOARCH=$(GOARCH) VERSION=$(VERSION) nfpm package --packager rpm --target featurebase.$(VERSION).$(GOARCH).rpm - + # We allow setting a custom docker-compose "project". Multiple of the # same docker-compose environment can exist simultaneously as long as # they use different projects (the project name is prepended to @@ -122,10 +124,15 @@ authclustertests: vendor PROJECT=$(PROJECT) ENABLE_AUTH=1 $(DOCKER_COMPOSE) -f internal/clustertests/docker-compose.yml run client1 CLUSTERTESTS_FB_ARGS=$(AUTH_ARGS) $(DOCKER_COMPOSE) -f internal/clustertests/docker-compose.yml down -# Install FeatureBase -install: +# Install FeatureBase and IDK +install: install-featurebase install-idk + +install-featurebase: $(GO) install -tags='$(BUILD_TAGS)' -ldflags $(LDFLAGS) $(FLAGS) ./cmd/featurebase +install-idk: + $(MAKE) -C ./idk install + # Build the lattice assets build-lattice: docker build -t lattice:build ./lattice @@ -187,6 +194,48 @@ docker-image: vendor --tag featurebase:$(VERSION) . @echo Created docker image: featurebase:$(VERSION) +docker-image-featurebase: vendor + docker build \ + --build-arg GO_VERSION=$(GO_VERSION) \ + --file Dockerfile-dax \ + --tag dax/featurebase . + +docker-image-featurebase-test: vendor + docker build \ + --build-arg GO_VERSION=$(GO_VERSION) \ + --file Dockerfile-clustertests \ + --tag dax/featurebase-test . + + +# build-for-quick builds a linux featurebase binary outside of docker +# (which is much faster for some reason), and places it in the .quick +# subdirectory. +build-for-quick: + GOOS=linux $(MAKE) build FLAGS="-o .quick/fb_linux" + +# docker-image-featurebase-quick uses a pre-built featurebase binary +# to quickly create a fresh docker image without needing to send the +# context of the featurebase top level directory. +docker-image-featurebase-quick: build-for-quick + docker build \ + --build-arg GO_VERSION=$(GO_VERSION) \ + --file Dockerfile-dax-quick ./.quick/ + + +docker-image-datagen: vendor + docker build --tag dax/datagen --file Dockerfile-datagen . + +ecr-push-featurebase: docker-login + docker tag dax/featurebase:latest $(AWS_ACCOUNTID).dkr.ecr.us-east-2.amazonaws.com/dax:latest + docker push $(AWS_ACCOUNTID).dkr.ecr.us-east-2.amazonaws.com/dax:latest + +ecr-push-datagen: docker-login + docker tag dax/datagen:latest $(AWS_ACCOUNTID).dkr.ecr.us-east-2.amazonaws.com/dax/datagen:latest + docker push $(AWS_ACCOUNTID).dkr.ecr.us-east-2.amazonaws.com/dax/datagen:latest + +docker-login: + aws ecr get-login-password --region us-east-2 | docker login --username AWS --password-stdin $(AWS_ACCOUNTID).dkr.ecr.us-east-2.amazonaws.com + # Create docker image (alias) docker: docker-image # alias diff --git a/api.go b/api.go index 396b5fa5a..6b6f95cfe 100644 --- a/api.go +++ b/api.go @@ -22,8 +22,10 @@ import ( "sync" "time" - "github.com/featurebasedb/featurebase/v3/disco" - "github.com/featurebasedb/featurebase/v3/rbf" + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/computer" + "github.com/molecula/featurebase/v3/disco" + "github.com/molecula/featurebase/v3/rbf" //"github.com/featurebasedb/featurebase/v3/pg" "github.com/featurebasedb/featurebase/v3/pql" @@ -51,6 +53,16 @@ type API struct { importWork chan importJob Serializer Serializer + + writeLogReader computer.WriteLogReader + writeLogWriter computer.WriteLogWriter + snapshotReadWriter computer.SnapshotReadWriter + + directiveWorkerPoolSize int + + // isComputeNode is set to true if this node is running as a DAX compute + // node. + isComputeNode bool } func (api *API) Holder() *Holder { @@ -77,10 +89,50 @@ func OptAPIImportWorkerPoolSize(size int) apiOption { } } +func OptAPIWriteLogReader(wlr computer.WriteLogReader) apiOption { + return func(a *API) error { + a.writeLogReader = wlr + return nil + } +} + +func OptAPIWriteLogWriter(wlw computer.WriteLogWriter) apiOption { + return func(a *API) error { + a.writeLogWriter = wlw + return nil + } +} + +func OptAPISnapshotter(snap computer.SnapshotReadWriter) apiOption { + return func(a *API) error { + a.snapshotReadWriter = snap + return nil + } +} + +func OptAPIDirectiveWorkerPoolSize(size int) apiOption { + return func(a *API) error { + a.directiveWorkerPoolSize = size + return nil + } +} + +func OptAPIIsComputeNode(is bool) apiOption { + return func(a *API) error { + a.isComputeNode = is + return nil + } +} + // NewAPI returns a new API instance. func NewAPI(opts ...apiOption) (*API, error) { api := &API{ importWorkerPoolSize: 2, + writeLogReader: computer.NewNopWriteLogReader(), + writeLogWriter: computer.NewNopWriteLogWriter(), + snapshotReadWriter: computer.NewNopSnapshotReadWriter(), + + directiveWorkerPoolSize: 2, } for _, opt := range opts { @@ -104,7 +156,7 @@ func NewAPI(opts ...apiOption) (*API, error) { return api, nil } -// Setter for API options. +// SetAPIOptions applies the given functional options to the API. func (api *API) SetAPIOptions(opts ...apiOption) error { for _, opt := range opts { err := opt(api) @@ -554,11 +606,20 @@ func (api *API) ImportRoaring(ctx context.Context, indexName, fieldName string, return errors.Wrap(err, "validating api method") } + api.server.logger.Debugf("ImportRoaring: %v %v %v", indexName, fieldName, shard) index, field, err := api.indexField(indexName, fieldName, shard) if index == nil || field == nil { return err } + // This node only handles the shard(s) that it owns. + if api.isComputeNode { + directive := api.holder.Directive() + if !shardInShards(dax.ShardNum(shard), directive.ComputeShards(dax.TableKey(index.Name()))) { + return errors.Errorf("import request shard is not supported (roaring): %d", shard) + } + } + if err = req.ValidateWithTimestamp(index.CreatedAt(), field.CreatedAt()); err != nil { return newPreconditionFailedError(err) } @@ -608,11 +669,66 @@ func (api *API) ImportRoaring(ctx context.Context, indexName, fieldName string, // Exit once all nodes are processed. if maxNode == len(nodes) { + if api.isComputeNode && !req.SuppressLog { + // Write the request to the write logger. + partition := disco.ShardToShardPartition(indexName, shard, disco.DefaultPartitionN) + msg := &computer.ImportRoaringMessage{ + Table: indexName, + Field: fieldName, + Partition: partition, + Shard: shard, + Clear: req.Clear, + Action: req.Action, + Block: req.Block, + UpdateExistence: req.UpdateExistence, + Views: req.Views, + } + + // Get the current version for shard. + version, err := api.getOrCreateShardVersion(ctx, indexName, shard) + if err != nil { + return errors.Wrap(err, "get or creating shard version") + } + + tkey := dax.TableKey(indexName) + qtid := tkey.QualifiedTableID() + partitionNum := dax.PartitionNum(partition) + shardNum := dax.ShardNum(shard) + + api.server.logger.Debugf("importroaring writing to writelogger: %+v, %[1]T len(msg.Views): %d, table: %s", api.writeLogWriter, len(msg.Views), msg.Table) + if err := api.writeLogWriter.WriteShard(ctx, qtid, partitionNum, shardNum, version, msg); err != nil { + return err + } + } + return qcx.Finish() } } } +func (api *API) getOrCreateShardVersion(ctx context.Context, indexName string, shard uint64) (int, error) { + tableName := dax.TableName(indexName) + shardNum := dax.ShardNum(shard) + + // Here we assume that indexName is the string encoding of QualifiedTableID. + qtid, err := dax.QualifiedTableIDFromKey(indexName) + if err != nil { + return -1, errors.Wrap(err, "decoding qtid from key (indexName)") + } + + version, found, err := api.holder.versionStore.ShardVersion(ctx, qtid, shardNum) + if err != nil { + return -1, errors.Wrap(err, "getting shard version") + } else if !found { + version = 0 + api.server.logger.Printf("could not find version for shard: %s, %d, so creating 0", tableName, shardNum) + if err := api.holder.versionStore.AddShards(ctx, qtid, dax.NewShard(shardNum, version)); err != nil { + return -1, errors.Wrap(err, "adding shard 0") + } + } + return version, nil +} + // DeleteField removes the named field from the named index. If the index is not // found, an error is returned. If the field is not found, it is ignored and no // action is taken. @@ -1050,6 +1166,22 @@ func (api *API) IndexInfo(ctx context.Context, name string) (*IndexInfo, error) return nil, ErrIndexNotFound } +// FieldInfo returns the same information as Schema(), but only for a single +// field. +func (api *API) FieldInfo(ctx context.Context, indexName, fieldName string) (*FieldInfo, error) { + idx, err := api.IndexInfo(ctx, indexName) + if err != nil { + return nil, err + } + + fld := idx.Field(fieldName) + if fld == nil { + return nil, ErrFieldNotFound + } + + return fld, nil +} + // ApplySchema takes the given schema and applies it across the // cluster (if remote is false), or just to this node (if remote is // true). This is designed for the use case of replicating a schema @@ -1183,6 +1315,7 @@ type ImportOptions struct { IgnoreKeyCheck bool Presorted bool fullySorted bool // format-aware sorting, internal use only please. + suppressLog bool // test Tx atomicity if > 0 SimPowerLossAfter int @@ -1216,6 +1349,13 @@ func OptImportOptionsPresorted(b bool) ImportOption { } } +func OptImportOptionsSuppressLog(b bool) ImportOption { + return func(o *ImportOptions) error { + o.suppressLog = b + return nil + } +} + var ErrAborted = fmt.Errorf("error: update was aborted") func (api *API) ImportAtomicRecord(ctx context.Context, qcx *Qcx, req *AtomicRecord, opts ...ImportOption) error { @@ -1306,10 +1446,70 @@ func (api *API) Import(ctx context.Context, qcx *Qcx, req *ImportRequest, opts . if err != nil { return errors.Wrap(err, "setting up import options") } + + ///////////////////////////////////////////////////////////////////////////// + // We build the ImportMessage here BEFORE the call to api.ImportWithTx(), + // because something in that method is modifying the values of req, so if we + // build ImportMessage after the call to api.ImportWithTx(), then the values + // that get logged are incorrect. An example I saw were RowIDS going from: + // shard 0 [1, 2, 3] + // shard 4 [0] + // + // to: + // shard 0 [1048577, 2097154, 3145731] + // shard 4 [1] + // + // which seem to be the offset in the field shard bitmap. + var partition int + var msg *computer.ImportMessage + if api.isComputeNode && !options.suppressLog { + partition = disco.ShardToShardPartition(req.Index, req.Shard, disco.DefaultPartitionN) + msg = &computer.ImportMessage{ + Table: req.Index, + Field: req.Field, + Partition: partition, + Shard: req.Shard, + RowIDs: make([]uint64, len(req.RowIDs)), + ColumnIDs: make([]uint64, len(req.ColumnIDs)), + RowKeys: make([]string, len(req.RowKeys)), + ColumnKeys: make([]string, len(req.ColumnKeys)), + Timestamps: make([]int64, len(req.Timestamps)), + Clear: req.Clear, + + IgnoreKeyCheck: options.IgnoreKeyCheck, + Presorted: options.Presorted, + } + copy(msg.RowIDs, req.RowIDs) + copy(msg.ColumnIDs, req.ColumnIDs) + copy(msg.RowKeys, req.RowKeys) + copy(msg.ColumnKeys, req.ColumnKeys) + copy(msg.Timestamps, req.Timestamps) + } + ///////////////////////////////////////////////////////////////////////////// + err = api.ImportWithTx(ctx, qcx, req, options) if err != nil { return err } + + if api.isComputeNode && !options.suppressLog { + // Get the current version for shard. + version, err := api.getOrCreateShardVersion(ctx, req.Index, req.Shard) + if err != nil { + return errors.Wrap(err, "get or creating shard version") + } + + tkey := dax.TableKey(req.Index) + qtid := tkey.QualifiedTableID() + partitionNum := dax.PartitionNum(partition) + shardNum := dax.ShardNum(req.Shard) + + // Write the request to the write logger. + if err := api.writeLogWriter.WriteShard(ctx, qtid, partitionNum, shardNum, version, msg); err != nil { + return err + } + } + return nil } @@ -1322,11 +1522,20 @@ func (api *API) ImportWithTx(ctx context.Context, qcx *Qcx, req *ImportRequest, return errors.Wrap(err, "validating api method") } + api.server.logger.Debugf("ImportWithTx: %v %v %v", req.Index, req.Field, req.Shard) idx, field, err := api.indexField(req.Index, req.Field, req.Shard) if err != nil { return errors.Wrap(err, "getting index and field") } + // This node only handles the shard(s) that it owns. + if api.isComputeNode { + directive := api.holder.Directive() + if !shardInShards(dax.ShardNum(req.Shard), directive.ComputeShards(dax.TableKey(idx.Name()))) { + return errors.Errorf("import request shard is not supported (with tx): %d", req.Shard) + } + } + if err := req.ValidateWithTimestamp(idx.CreatedAt(), field.CreatedAt()); err != nil { return errors.Wrap(err, "validating import value request") } @@ -1450,7 +1659,8 @@ func (api *API) ImportRoaringShard(ctx context.Context, indexName string, shard defer finisher(&err1) if !req.Remote { - return errors.New("forwarding unimplemented on this endpoint") + err1 = errors.New("forwarding unimplemented on this endpoint") + return err1 } for _, viewUpdate := range req.Views { @@ -1514,6 +1724,42 @@ func (api *API) ImportRoaringShard(ctx context.Context, indexName string, shard } } + if api.isComputeNode && !req.SuppressLog { + partition := disco.ShardToShardPartition(indexName, shard, disco.DefaultPartitionN) + msg := &computer.ImportRoaringShardMessage{ + Table: indexName, + Partition: partition, + Shard: shard, + Views: make([]computer.RoaringUpdate, len(req.Views)), + } + for i, view := range req.Views { + msg.Views[i] = computer.RoaringUpdate{ + Field: view.Field, + View: view.View, + Clear: view.Clear, + Set: view.Set, + ClearRecords: view.ClearRecords, + } + } + // Get the current version for shard. + version, err := api.getOrCreateShardVersion(ctx, indexName, shard) + if err != nil { + err1 = errors.Wrap(err, "get or creating shard version") + return err1 + } + tkey := dax.TableKey(indexName) + qtid := tkey.QualifiedTableID() + partitionNum := dax.PartitionNum(partition) + shardNum := dax.ShardNum(shard) + + api.server.logger.Debugf("importroaringshard writing shard to writelogger: %+v, len(msg.Views): %d, table: %s", api.writeLogWriter, len(msg.Views), msg.Table) + + if err := api.writeLogWriter.WriteShard(ctx, qtid, partitionNum, shardNum, version, msg); err != nil { + err1 = errors.Wrap(err, "writing import-roaring-shard to writelogger") + return err1 + } + } + return nil } @@ -1549,7 +1795,64 @@ func (api *API) ImportValue(ctx context.Context, qcx *Qcx, req *ImportValueReque if err != nil { return errors.Wrap(err, "setting up import options") } - return api.ImportValueWithTx(ctx, qcx, req, options) + + ///////////////////////////////////////////////////////////////////////////// + // We build the ImportValueMessage here BEFORE the call to + // api.ImportValueWithTx() because we don't trust that req doesn't get + // changed out from under us. See the similar comment in the API.Import() + // method above. + var partition int + var msg *computer.ImportValueMessage + if api.isComputeNode && !options.suppressLog { + partition = disco.ShardToShardPartition(req.Index, req.Shard, disco.DefaultPartitionN) + msg = &computer.ImportValueMessage{ + Table: req.Index, + Field: req.Field, + Partition: partition, + Shard: req.Shard, + ColumnIDs: make([]uint64, len(req.ColumnIDs)), + ColumnKeys: make([]string, len(req.ColumnKeys)), + Values: make([]int64, len(req.Values)), + FloatValues: make([]float64, len(req.FloatValues)), + TimestampValues: make([]time.Time, len(req.TimestampValues)), + StringValues: make([]string, len(req.StringValues)), + Clear: req.Clear, + + IgnoreKeyCheck: options.IgnoreKeyCheck, + Presorted: options.Presorted, + } + copy(msg.ColumnIDs, req.ColumnIDs) + copy(msg.ColumnKeys, req.ColumnKeys) + copy(msg.Values, req.Values) + copy(msg.FloatValues, req.FloatValues) + copy(msg.TimestampValues, req.TimestampValues) + copy(msg.StringValues, req.StringValues) + } + ///////////////////////////////////////////////////////////////////////////// + + if err := api.ImportValueWithTx(ctx, qcx, req, options); err != nil { + return errors.Wrap(err, "importing value with tx") + } + + if api.isComputeNode && !options.suppressLog { + // Get the current version for shard. + version, err := api.getOrCreateShardVersion(ctx, req.Index, req.Shard) + if err != nil { + return errors.Wrap(err, "get or creating shard version") + } + + tkey := dax.TableKey(req.Index) + qtid := tkey.QualifiedTableID() + partitionNum := dax.PartitionNum(partition) + shardNum := dax.ShardNum(req.Shard) + + // Write the request to the write logger. + if err := api.writeLogWriter.WriteShard(ctx, qtid, partitionNum, shardNum, version, msg); err != nil { + return errors.Wrap(err, "writing shard to write logger") + } + } + + return nil } // ImportValueWithTx bulk imports values into a particular field. @@ -1570,19 +1873,24 @@ func (api *API) ImportValueWithTx(ctx context.Context, qcx *Qcx, req *ImportValu return nil } + api.server.logger.Debugf("ImportValueWithTx: %v %v %v", req.Index, req.Field, req.Shard) idx, field, err := api.indexField(req.Index, req.Field, req.Shard) if err != nil { return errors.Wrap(err, fmt.Sprintf("getting index '%v' and field '%v'; shard=%v", req.Index, req.Field, req.Shard)) } + // This node only handles the shard(s) that it owns. + if api.isComputeNode { + directive := api.holder.Directive() + if !shardInShards(dax.ShardNum(req.Shard), directive.ComputeShards(dax.TableKey(idx.Name()))) { + return errors.Errorf("import request shard is not supported (value with tx): %d", req.Shard) + } + } + if err := req.ValidateWithTimestamp(idx.CreatedAt(), field.CreatedAt()); err != nil { return errors.Wrap(err, "validating import value request") } - idx, field, err = api.indexField(req.Index, req.Field, req.Shard) - if err != nil { - return errors.Wrap(err, "getting index and field") - } span.LogKV( "index", req.Index, "field", req.Field) @@ -1821,8 +2129,6 @@ func (api *API) validateShardOwnership(indexName string, shard uint64) error { } func (api *API) indexField(indexName string, fieldName string, shard uint64) (*Index, *Field, error) { - api.server.logger.Debugf("importing: %v %v %v", indexName, fieldName, shard) - // Find the Index. index := api.holder.Index(indexName) if index == nil { @@ -2270,7 +2576,7 @@ func (api *API) TranslateFieldDB(ctx context.Context, indexName, fieldName strin return err } -// RestoreShard +// RestoreShard is used by the restore tool to restore previously backed up data. This call is specific to RBF data for a shard. func (api *API) RestoreShard(ctx context.Context, indexName string, shard uint64, rd io.Reader) error { snap := api.cluster.NewSnapshot() if !snap.OwnsShard(api.server.nodeID, indexName, shard) { @@ -2760,6 +3066,144 @@ func (api *API) RBFDebugInfo() map[string]*rbf.DebugInfo { return infos } +func (api *API) Directive(ctx context.Context, d *dax.Directive) error { + return api.ApplyDirective(ctx, d) +} + +// SnapshotShardData triggers the node to perform a shard snapshot based on the +// provided SnapshotShardDataRequest. +func (api *API) SnapshotShardData(ctx context.Context, req *dax.SnapshotShardDataRequest) error { + qtid := req.TableKey.QualifiedTableID() + + // Confirm that this node is currently responsible for table/shard/fromVersion. + var version int + if v, ok, err := api.holder.versionStore.ShardVersion(ctx, qtid, req.ShardNum); err != nil { + return err + } else if !ok { + return errors.Errorf("shard not managed by this node: %s, %d", req.TableKey, req.ShardNum) + } else if v != req.FromVersion { + return errors.Errorf("shard managed by this node is at version: %d, not: %d", v, req.FromVersion) + } else { + version = v + } + + partition := disco.ShardToShardPartition(string(req.TableKey), uint64(req.ShardNum), disco.DefaultPartitionN) + partitionNum := dax.PartitionNum(partition) + + // Create the snapshot for the current version. + rc, err := api.IndexShardSnapshot(ctx, string(req.TableKey), uint64(req.ShardNum)) + if err != nil { + return errors.Wrap(err, "getting index/shard readcloser") + } + + // The following closes rc, the ReadCloser. + if err := api.snapshotReadWriter.WriteShardData(ctx, qtid, partitionNum, req.ShardNum, version, rc); err != nil { + return errors.Wrap(err, "snapshotting shard data") + } + + // Increment the version of the shard managed by this node. + if err := api.holder.versionStore.AddShards(ctx, qtid, + dax.NewShard(req.ShardNum, req.ToVersion), + ); err != nil { + return errors.Wrap(err, "incrementing shard version locally") + } + + // Update the cached directive on the holder. + api.holder.SetDirective(&req.Directive) + + // Finally, delete the log file for the previous version. + return api.writeLogWriter.DeleteShard(ctx, qtid, partitionNum, req.ShardNum, req.FromVersion) +} + +// SnapshotTableKeys triggers the node to perform a table keys snapshot based on +// the provided SnapshotTableKeysRequest. +func (api *API) SnapshotTableKeys(ctx context.Context, req *dax.SnapshotTableKeysRequest) error { + // If the index is not keyed, no-op on snapshotting its keys. + if idx, err := api.Index(ctx, string(req.TableKey)); err != nil { + return newNotFoundError(ErrIndexNotFound, string(req.TableKey)) + } else if !idx.Keys() { + return nil + } + + qtid := req.TableKey.QualifiedTableID() + + // Confirm that this node is currently responsible for table/partition/fromVersion. + var version int + if v, ok, err := api.holder.versionStore.PartitionVersion(ctx, qtid, req.PartitionNum); err != nil { + return err + } else if !ok { + return errors.Errorf("partition not managed by this node: %s, %d", req.TableKey, req.PartitionNum) + } else if v != req.FromVersion { + return errors.Errorf("partition managed by this node is at version: %d, not: %d", v, req.FromVersion) + } else { + version = v + } + + // Create the snapshot for the current version. + wrTo, err := api.TranslateData(ctx, string(req.TableKey), int(req.PartitionNum)) + if err != nil { + return errors.Wrapf(err, "getting index/partition writeto: %s/%d", req.TableKey, req.PartitionNum) + } + + if err := api.snapshotReadWriter.WriteTableKeys(ctx, qtid, req.PartitionNum, version, wrTo); err != nil { + return errors.Wrap(err, "snapshotting table keys") + } + + // Increment the version of the partition managed by this node. + if err := api.holder.versionStore.AddPartitions(ctx, qtid, + dax.NewPartition(req.PartitionNum, req.ToVersion), + ); err != nil { + return errors.Wrap(err, "incrementing partition version locally") + } + + // Update the cached directive on the holder. + api.holder.SetDirective(&req.Directive) + + // Finally, delete the log file for the previous version. + return api.writeLogWriter.DeleteTableKeys(ctx, qtid, req.PartitionNum, req.FromVersion) +} + +// SnapshotFieldKeys triggers the node to perform a field keys snapshot based on +// the provided SnapshotFieldKeysRequest. +func (api *API) SnapshotFieldKeys(ctx context.Context, req *dax.SnapshotFieldKeysRequest) error { + qtid := req.TableKey.QualifiedTableID() + + // Confirm that this node is currently responsible for table/field/fromVersion. + var version int + if v, ok, err := api.holder.versionStore.FieldVersion(ctx, qtid, req.Field); err != nil { + return err + } else if !ok { + return errors.Errorf("field not managed by this node: %s, %s", req.TableKey, req.Field) + } else if v != req.FromVersion { + return errors.Errorf("field managed by this node is at version: %d, not: %d", v, req.FromVersion) + } else { + version = v + } + + // Create the snapshot for the current version. + wrTo, err := api.FieldTranslateData(ctx, string(req.TableKey), string(req.Field)) + if err != nil { + return errors.Wrap(err, "getting index/field writeto") + } + + if err := api.snapshotReadWriter.WriteFieldKeys(ctx, qtid, req.Field, version, wrTo); err != nil { + return errors.Wrap(err, "snapshotting field keys") + } + + // Increment the version of the field managed by this node. + if err := api.holder.versionStore.AddFields(ctx, qtid, + dax.NewFieldVersion(req.Field, req.ToVersion), + ); err != nil { + return errors.Wrap(err, "incrementing field version locally") + } + + // Update the cached directive on the holder. + api.holder.SetDirective(&req.Directive) + + // Finally, delete the log file for the previous version. + return api.writeLogWriter.DeleteFieldKeys(ctx, qtid, req.Field, req.FromVersion) +} + type serverInfo struct { ShardWidth uint64 `json:"shardWidth"` ReplicaN int `json:"replicaN"` @@ -2880,19 +3324,39 @@ var methodsNormal = map[apiMethod]struct{}{ apiMutexCheck: {}, } +func shardInShards(i dax.ShardNum, s dax.Shards) bool { + for _, o := range s { + if i == o.Num { + return true + } + } + return false +} + // SchemaAPI is a subset of the API methods which have to do with schema. This // interface was introduced in order to remove, from the sql3 package, the // pointer to API, and instead use this interface. In the current FeatureBase, -// this interface can be implemented directly with API. But in an implementation -// for DAX, for example, we might want something else servicing the -// schema-related calls to the SchemaAPI. +// this interface can be implemented directly with API (well, not directly, but +// with FeatureBaseSchemaAPI, which is a wrapper around API). But in an +// implementation for DAX, for example, we might want something else servicing +// the schema-related calls to the SchemaAPI. type SchemaAPI interface { CreateIndexAndFields(ctx context.Context, indexName string, options IndexOptions, fields []CreateFieldObj) error CreateField(ctx context.Context, indexName string, fieldName string, opts ...FieldOption) (*Field, error) DeleteField(ctx context.Context, indexName string, fieldName string) error DeleteIndex(ctx context.Context, indexName string) error - IndexInfo(ctx context.Context, indexName string) (*IndexInfo, error) + + // Schema returns the list of tables and fields. While it might make sense + // to have this as part of the SchemaInfoAPI interface instead of here, it's + // never used by consumers of that interface. Schema(ctx context.Context, withViews bool) ([]*IndexInfo, error) + + SchemaInfoAPI +} + +type SchemaInfoAPI interface { + IndexInfo(ctx context.Context, indexName string) (*IndexInfo, error) + FieldInfo(ctx context.Context, indexName, fieldName string) (*FieldInfo, error) } type ClusterNode struct { @@ -2936,6 +3400,28 @@ type QueryAPI interface { Query(ctx context.Context, req *QueryRequest) (QueryResponse, error) } +// Ensure type implements interface. +var _ ComputeAPI = (*NopComputeAPI)(nil) + +// NopComputeAPI is a no-op implementation of the ComputeAPI interface. +type NopComputeAPI struct{} + +func NewNopComputeAPI() *NopComputeAPI { + return &NopComputeAPI{} +} + +func (c *NopComputeAPI) Import(ctx context.Context, qcx *Qcx, req *ImportRequest, opts ...ImportOption) error { + return nil +} +func (c *NopComputeAPI) ImportValue(ctx context.Context, qcx *Qcx, req *ImportValueRequest, opts ...ImportOption) error { + return nil +} + +func (c *NopComputeAPI) Txf() *TxFactory { return nil } + +// Ensure type implements interface. +var _ SchemaAPI = (*FeatureBaseSchemaAPI)(nil) + // FeatureBaseSchemaAPI is a wrapper around pilosa.API. It implements the // SchemaAPI interface with methods which are not a part of pilosa.API. type FeatureBaseSchemaAPI struct { diff --git a/api_directive.go b/api_directive.go new file mode 100644 index 000000000..a907ea442 --- /dev/null +++ b/api_directive.go @@ -0,0 +1,983 @@ +// Copyright 2021 Molecula Corp. All rights reserved. +package pilosa + +import ( + "context" + "io" + "log" + "sync" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/computer" + "github.com/molecula/featurebase/v3/disco" + "github.com/pkg/errors" +) + +// ApplyDirective applies a Directive received, from the Controller, at the +// /directive endpoint. +func (api *API) ApplyDirective(ctx context.Context, d *dax.Directive) error { + // Get the current directive for comparison. + previousDirective := api.holder.Directive() + + // Check that incoming version is newer. + // Note: 0 is an invalid Directive version. This decision was made because + // previousDirective is not a pointer to a directive, but a concrete + // Directive. Which means we can't check for nil, and by default it has a + // version of 0. So in order to ensure the version has increased, we need to + // require that incoming directive versions are greater than 0. + if d.Version == 0 { + return errors.Errorf("directive version cannot be 0") + } else if previousDirective.Version >= d.Version { + return errors.Errorf("directive version mismatch, got %d, but already have %d", d.Version, previousDirective.Version) + } + + // Handle the operations based on the directive method. + switch d.Method { + case dax.DirectiveMethodDiff: + // pass: normal operation + + case dax.DirectiveMethodReset: + // Delete all tables. + if err := api.deleteAllIndexes(ctx); err != nil { + return errors.Wrap(err, "deleting all indexes") + } + + // Set previousDirective to empty so the diff handles everything as new. + previousDirective = dax.Directive{} + + case dax.DirectiveMethodSnapshot: + // TODO(tlt): this was the existing logic, but we should really diff the + // directive and ensure that overwriting the value in the cache doesn't + // have a negative effect. + api.holder.SetDirective(d) + return nil + + default: + return errors.Errorf("invalid directive method: %s", d.Method) + } + + // Cache this directive as the latest applied. There is functionality within + // the "enactDirective" stage of ApplyDirective which validates against this + // cached Directive, so it's important that it be set before calling + // enactDirective(). An example: when loading partition data from the + // WriteLogger, there are validations to ensure that the partition being + // loaded is meant to be handled by this node; that validation is done + // against the cached Directive. + // TODO(tlt): despite what this comment says, this logic is not sound; we + // shouldn't be setting the directive until enactiveDirective() succeeds. + api.holder.SetDirective(d) + + return api.enactDirective(ctx, &previousDirective, d) +} + +// deleteAllIndexes deletes all indexes handled by this node. +func (api *API) deleteAllIndexes(ctx context.Context) error { + indexes, err := api.Schema(ctx, false) + if err != nil { + return errors.Wrap(err, "getting schema") + } + + for i := range indexes { + if err := api.DeleteIndex(ctx, indexes[i].Name); err != nil { + return errors.Wrapf(err, "deleting index: %s", indexes[i].Name) + } + } + + return nil +} + +// directiveJobType allows us to switch on jobType in the directiveWorker in +// order to use a single worker pool for all job types (as opposed to having a +// separate worker pool for each job type). +type directiveJobType interface { + // We have this method just to prevent *any* struct from implementing this + // interface automatically. But, interestingly enough, we don't actually + // have to have this method on the implementation because we embed the + // interface. + isJobType() bool +} + +type directiveJobTableKeys struct { + directiveJobType + idx *Index + tkey dax.TableKey + partition dax.Partition +} + +type directiveJobFieldKeys struct { + directiveJobType + tkey dax.TableKey + field dax.FieldVersion +} + +type directiveJobShards struct { + directiveJobType + tkey dax.TableKey + shard dax.Shard +} + +// directiveWorker is a worker in a worker pool which handles portions of a +// directive. Multiple instances of directiveWorker run in goroutines in order +// to load data from snapshotter and writelogger concurrently. Note: unlike the +// api.ingestWorkerPool, of which one pool is always running, the +// directiveWorker pool is only running during the life of the +// api.ApplyDirective call. Technically, this means that multiple +// directiveWorker pools could be active at the same time, but we should never +// be running more than once instance of ApplyDirective concurrently. +func (api *API) directiveWorker(ctx context.Context, jobs <-chan directiveJobType, errs chan<- error) { + for j := range jobs { + switch job := j.(type) { + case directiveJobTableKeys: + if err := api.loadTableKeys(ctx, job.idx, job.tkey, job.partition); err != nil { + errs <- errors.Wrapf(err, "loading table keys: %s, %s", job.tkey, job.partition) + } + case directiveJobFieldKeys: + if err := api.loadFieldKeys(ctx, job.tkey, job.field); err != nil { + errs <- errors.Wrapf(err, "loading field keys: %s, %s", job.tkey, job.field) + } + case directiveJobShards: + if err := api.loadShard(ctx, job.tkey, job.shard); err != nil { + errs <- errors.Wrapf(err, "loading shard: %s, %s", job.tkey, job.shard) + } + default: + errs <- errors.Errorf("unsupported job type: %T %[1]v", job) + } + + select { + case <-ctx.Done(): + return + default: + // continue pulling jobs off the channel + } + } +} + +func (api *API) enactDirective(ctx context.Context, fromD, toD *dax.Directive) error { + // enactTables is called before the jobs that run in the worker pool because + // it probably makes sense to apply the schema before trying to load data + // concurrently. + if err := api.enactTables(ctx, fromD, toD); err != nil { + return errors.Wrap(err, "enactTables") + } + + // The following types use a shared pool of workers to run each + // directiveJobType. + + var wg sync.WaitGroup + + // open job channel + jobs := make(chan directiveJobType, api.directiveWorkerPoolSize) + errs := make(chan error) + done := make(chan struct{}) + + // Spin up n workers in goroutines that pull jobs from the jobs channel. + for i := 0; i < api.directiveWorkerPoolSize; i++ { + wg.Add(1) + go func() { + api.directiveWorker(ctx, jobs, errs) + defer wg.Done() + }() + } + + // Wait for the WaitGroup counter to reach 0. When it has, indicate that + // we're done processing all jobs by closing the done channel. + go func() { + wg.Wait() + close(done) + }() + + // Run through all the "enact" methods. These push jobs onto the jobs + // channel. Once all the jobs have been queued to the channel, we close the + // jobs channel. This allows the directiveWorkers to exit out of the + // function, which will then decrement the WaitGroup counter. + go func() { + api.pushJobsTableKeys(ctx, jobs, fromD, toD) + api.pushJobsFieldKeys(ctx, jobs, fromD, toD) + api.pushJobsShards(ctx, jobs, fromD, toD) + close(jobs) + }() + + // Keep running until we get an error or until the done channel is closed. + // Note: the code is written such that only non-nil errors are pushed to the + // errs channel. + for { + select { + case err := <-errs: + return err + case <-done: + return nil + } + } +} + +func (api *API) enactTables(ctx context.Context, fromD, toD *dax.Directive) error { + currentIndexes := api.holder.Indexes() + + // Make a list of indexes that currently exist (from). + from := make(dax.TableKeys, 0, len(currentIndexes)) + for _, idx := range currentIndexes { + qtid, err := dax.QualifiedTableIDFromKey(idx.Name()) + if err != nil { + return errors.Wrap(err, "converting index name to qualified table id") + } + from = append(from, qtid.Key()) + } + + // TODO sanity check holder against fromD. We're getting existing + // indexes from holder, but in theory fromD should be + // identical. If we have an error in our directive-caching logic + // (it has happened before (just now, in fact!) and we'd be + // foolish to think it won't happen again), or we have schema + // mutations that are not going through the directive path, we + // could potentially catch them here. + + // Make a list of tables that are in the directive (to) along with a map of + // tableKey to table (m). + m := make(map[dax.TableKey]*dax.QualifiedTable, len(toD.Tables)) + to := make(dax.TableKeys, 0, len(toD.Tables)) + for _, t := range toD.Tables { + m[t.Key()] = t + to = append(to, t.Key()) + } + + sc := newSliceComparer(from, to) + + // Remove all indexes that are no longer part of the directive. + for _, tkey := range sc.removed() { + idx := string(tkey) + if err := api.holder.deleteIndex(idx); err != nil { + return errors.Wrapf(err, "deleting index: %s", tkey) + } + } + + // Put partitions into a map by table. + partitionMap := toD.TranslatePartitionsMap() + + // Add all indexes that weren't previously (but now are) a part of the + // directive. + for _, tkey := range sc.added() { + if qtbl, found := m[tkey]; !found { + return errors.Errorf("table '%s' was not in map", tkey) + } else if err := api.createTableAndFields(qtbl, partitionMap[tkey]); err != nil { + return err + } + } + + // Check fields on all indexes present in both from and to. + for _, tkey := range sc.same() { + if err := api.enactFieldsForTable(ctx, tkey, fromD, toD); err != nil { + return errors.Wrapf(err, "enacting fields for table: '%s'", tkey) + } + } + + return nil +} + +func (api *API) enactFieldsForTable(ctx context.Context, tkey dax.TableKey, fromD, toD *dax.Directive) error { + qtid := tkey.QualifiedTableID() + + fromT, err := fromD.Table(qtid) + if err != nil { + return errors.Wrap(err, "getting from table") + } + toT, err := toD.Table(qtid) + if err != nil { + return errors.Wrap(err, "getting to table") + } + + // Get the index for tkey. + idx := api.holder.Index(string(tkey)) + if idx == nil { + return errors.Errorf("index not found: %s", tkey) + } + + sc := newSliceComparer(fromT.FieldNames(), toT.FieldNames()) + + // Add fields new to toT. + for _, fldName := range sc.added() { + if field, found := toT.Field(fldName); !found { + return dax.NewErrFieldDoesNotExist(fldName) + } else if err := createField(idx, field); err != nil { + return errors.Wrapf(err, "creating field: %s/%s", tkey, fldName) + } + } + + // Remove fields which don't exist in toT. + for _, fldName := range sc.removed() { + if err := api.DeleteField(ctx, string(tkey), string(fldName)); err != nil { + return errors.Wrapf(err, "deleting field: %s/%s", tkey, fldName) + } + } + + // // Update any field options which have changed for existing fields. + // for _, fldName := range sc.same() { + // // handle changed field options?? + // } + + return nil +} + +func (api *API) pushJobsTableKeys(ctx context.Context, jobs chan<- directiveJobType, fromD, toD *dax.Directive) { + toPartitionsMap := toD.TranslatePartitionsMap() + + // Get the diff between from/to directive.partitions. + partComp := newPartitionsComparer(fromD.TranslatePartitionsMap(), toPartitionsMap) + + // Loop over the partition map and load from WriteLogger. + for tkey, partitions := range partComp.added() { + // Get index in order to find the translate stores (by partition) for + // the table. + idx := api.holder.Index(string(tkey)) + if idx == nil { + log.Printf("index not found in holder: %s", tkey) + continue + } + + // Update the cached version of translate partitions that we keep on the + // Index. + idx.SetTranslatePartitions(toPartitionsMap[tkey]) + + for _, partition := range partitions { + jobs <- directiveJobTableKeys{ + idx: idx, + tkey: tkey, + partition: partition, + } + } + } +} + +func (api *API) loadTableKeys(ctx context.Context, idx *Index, tkey dax.TableKey, partition dax.Partition) error { + qtid := tkey.QualifiedTableID() + + // Load the previous snapshot. Version 0 doesn't have a snapshot + // file; it only has log entries. + if partition.Version > 0 { + // Load partition snapshot: version - 1 + previousVersion := partition.Version - 1 + rc, err := api.snapshotReadWriter.ReadTableKeys(ctx, qtid, partition.Num, previousVersion) + if err != nil { + return errors.Wrap(err, "reading table keys snapshot") + } + defer rc.Close() + + if err := api.TranslateIndexDB(ctx, string(tkey), int(partition.Num), rc); err != nil { + return errors.Wrap(err, "restoring table keys") + + } + } + + if err := func() error { + store := idx.TranslateStore(int(partition.Num)) + + reader := api.writeLogReader.TableKeyReader(ctx, qtid, partition.Num, partition.Version) + if err := reader.Open(); err != nil { + // TODO: this log can be confusing because on a create + // table, there is no log file yet, so an error is expected. + // Instead of swallowing this error, we need to check the + // error code and handle it differently. This means the + // writelogger will need to return an error indicating that + // the log file does not exist, but that that is expected. + // log.Printf("could not open log file for table: %s, partition: %d: version: %d, err: %s", table, partition.Num, partition.Version, err) + return nil + } + defer reader.Close() + + for msg, err := reader.Read(); err != io.EOF; msg, err = reader.Read() { + if err != nil { + return errors.Wrap(err, "reading from log reader") + } + for key, id := range msg.StringToID { + if err := store.ForceSet(id, key); err != nil { + return errors.Wrapf(err, "forcing set id, key: %d, %s", id, key) + } + } + } + + return nil + }(); err != nil { + return err + } + + // Set the table/partition/version in the holder. + if err := api.holder.versionStore.AddPartitions(ctx, qtid, partition); err != nil { + return errors.Wrap(err, "adding partition to sharder") + } + + return nil +} + +func (api *API) pushJobsFieldKeys(ctx context.Context, jobs chan<- directiveJobType, fromD, toD *dax.Directive) { + // Get the diff between from/to directive.fields. + fieldComp := newFieldsComparer(fromD.TranslateFieldsMap(), toD.TranslateFieldsMap()) + + // Loop over the field map and load from WriteLogger. + for tkey, fields := range fieldComp.added() { + for _, field := range fields { + jobs <- directiveJobFieldKeys{ + tkey: tkey, + field: field, + } + } + } +} + +func (api *API) loadFieldKeys(ctx context.Context, tkey dax.TableKey, field dax.FieldVersion) error { + qtid := tkey.QualifiedTableID() + + // Load the previous snapshot. Version 0 doesn't have a snapshot + // file; it only has log entries. + if field.Version > 0 { + // Load field snapshot: version - 1 + previousVersion := field.Version - 1 + rc, err := api.snapshotReadWriter.ReadFieldKeys(ctx, qtid, field.Name, previousVersion) + if err != nil { + return errors.Wrap(err, "reading field keys snapshot") + } + defer rc.Close() + + if err := api.TranslateFieldDB(ctx, string(tkey), string(field.Name), rc); err != nil { + return errors.Wrap(err, "restoring field keys") + } + } + + if err := func() error { + // Get field in order to find the translate store. + fld := api.holder.Field(string(tkey), string(field.Name)) + if fld == nil { + log.Printf("field not found in holder: %s", field.Name) + return nil + } + store := fld.TranslateStore() + + reader := api.writeLogReader.FieldKeyReader(ctx, qtid, field.Name, field.Version) + if err := reader.Open(); err != nil { + // TODO: this log can be confusing because on a create + // table, there is no log file yet, so an error is expected. + // Instead of swallowing this error, we need to check the + // error code and handle it differently. This means the + // writelogger will need to return an error indicating that + // the log file does not exist, but that that is expected. + // log.Printf("could not open log file for table: %s, field: %s: version: %d, err: %s", table, field.Name, field.Version, err) + return nil + } + defer reader.Close() + + for msg, err := reader.Read(); err != io.EOF; msg, err = reader.Read() { + if err != nil { + return errors.Wrap(err, "reading from log reader") + } + for key, id := range msg.StringToID { + if err := store.ForceSet(id, key); err != nil { + return errors.Wrapf(err, "forcing set id, key: %d, %s", id, key) + } + } + } + + return nil + }(); err != nil { + return err + } + + // Set the table/field/version in the holder. + if err := api.holder.versionStore.AddFields(ctx, qtid, field); err != nil { + return errors.Wrap(err, "adding field to sharder") + } + + return nil +} + +func (api *API) pushJobsShards(ctx context.Context, jobs chan<- directiveJobType, fromD, toD *dax.Directive) { + // Put shards into a map by table. + shardMap := toD.ComputeShardsMap() + + // Get the diff between from/to directive shards. + shardComp := newShardsComparer(fromD.ComputeShardsMap(), shardMap) + + // Loop over the shard map and load from WriteLogger. + for tkey, shards := range shardComp.added() { + for _, shard := range shards { + jobs <- directiveJobShards{ + tkey: tkey, + shard: shard, + } + } + } +} + +func (api *API) loadShard(ctx context.Context, tkey dax.TableKey, shard dax.Shard) error { + qtid := tkey.QualifiedTableID() + + partition := disco.ShardToShardPartition(string(tkey), uint64(shard.Num), disco.DefaultPartitionN) + partitionNum := dax.PartitionNum(partition) + + // Load the previous snapshot. Version 0 doesn't have a snapshot + // file; it only has log entries. + if shard.Version > 0 { + // Load shard snapshot: version - 1 + previousVersion := shard.Version - 1 + rc, err := api.snapshotReadWriter.ReadShardData(ctx, qtid, partitionNum, shard.Num, previousVersion) + if err != nil { + return errors.Wrap(err, "reading shard data snapshot") + } + + if err := api.RestoreShard(ctx, string(tkey), uint64(shard.Num), rc); err != nil { + return errors.Wrap(err, "restoring shard data") + } + } + + // WriteLog reader. + if err := func() error { + reader := api.writeLogReader.ShardReader(ctx, qtid, partitionNum, shard.Num, shard.Version) + if err := reader.Open(); err != nil { + // TODO: this log can be confusing because on a create + // table, there is no log file yet, so an error is expected. + // Instead of swallowing this error, we need to check the + // error code and handle it differently. This means the + // writelogger will need to return an error indicating that + // the log file does not exist, but that that is expected. + // log.Printf("could not open log file for table: %s, partition: %d: version: %d, shard: %d, err: %s", table, partition, shard.Version, shard.Num, err) + return nil + } + defer reader.Close() + + for logMsg, err := reader.Read(); err != io.EOF; logMsg, err = reader.Read() { + if err != nil { + return errors.Wrap(err, "reading from log reader") + } + switch msg := logMsg.(type) { + case *computer.ImportRoaringMessage: + req := &ImportRoaringRequest{ + Clear: msg.Clear, + Action: msg.Action, + Block: msg.Block, + Views: msg.Views, + UpdateExistence: msg.UpdateExistence, + SuppressLog: true, + } + if err := api.ImportRoaring(ctx, msg.Table, msg.Field, msg.Shard, true, req); err != nil { + return errors.Wrapf(err, "import roaring, table: %s, field: %s, shard: %d", msg.Table, msg.Field, msg.Shard) + } + + case *computer.ImportMessage: + req := &ImportRequest{ + Index: msg.Table, + Field: msg.Field, + Shard: msg.Shard, + RowIDs: msg.RowIDs, + ColumnIDs: msg.ColumnIDs, + RowKeys: msg.RowKeys, + ColumnKeys: msg.ColumnKeys, + Timestamps: msg.Timestamps, + Clear: msg.Clear, + } + + qcx := api.Txf().NewQcx() + defer qcx.Abort() + + opts := []ImportOption{ + OptImportOptionsClear(msg.Clear), + OptImportOptionsIgnoreKeyCheck(msg.IgnoreKeyCheck), + OptImportOptionsPresorted(msg.Presorted), + OptImportOptionsSuppressLog(true), + } + if err := api.Import(ctx, qcx, req, opts...); err != nil { + return errors.Wrapf(err, "import, table: %s, field: %s, shard: %d", msg.Table, msg.Field, msg.Shard) + } + + case *computer.ImportValueMessage: + req := &ImportValueRequest{ + Index: msg.Table, + Field: msg.Field, + Shard: msg.Shard, + ColumnIDs: msg.ColumnIDs, + ColumnKeys: msg.ColumnKeys, + Values: msg.Values, + FloatValues: msg.FloatValues, + TimestampValues: msg.TimestampValues, + StringValues: msg.StringValues, + Clear: msg.Clear, + } + + qcx := api.Txf().NewQcx() + defer qcx.Abort() + + opts := []ImportOption{ + OptImportOptionsClear(msg.Clear), + OptImportOptionsIgnoreKeyCheck(msg.IgnoreKeyCheck), + OptImportOptionsPresorted(msg.Presorted), + OptImportOptionsSuppressLog(true), + } + if err := api.ImportValue(ctx, qcx, req, opts...); err != nil { + return errors.Wrapf(err, "import value, table: %s, field: %s, shard: %d", msg.Table, msg.Field, msg.Shard) + } + case *computer.ImportRoaringShardMessage: + req := &ImportRoaringShardRequest{ + Remote: true, + Views: make([]RoaringUpdate, len(msg.Views)), + SuppressLog: true, + } + for i, view := range msg.Views { + req.Views[i] = RoaringUpdate{ + Field: view.Field, + View: view.View, + Clear: view.Clear, + Set: view.Set, + ClearRecords: view.ClearRecords, + } + } + if err := api.ImportRoaringShard(ctx, msg.Table, msg.Shard, req); err != nil { + return errors.Wrapf(err, "import roaring shard table: %s, shard: %d", msg.Table, msg.Shard) + } + } + } + + return nil + }(); err != nil { + return err + } + + // Set the table/shard/version in the holder. + if err := api.holder.versionStore.AddShards(ctx, qtid, shard); err != nil { + return errors.Wrap(err, "adding shard to sharder") + } + + return nil +} + +////////////////////////////////////////////////////////////// + +// sliceComparer is used to compare the differences between two slices of comparables. +type sliceComparer[K comparable] struct { + from []K + to []K +} + +func newSliceComparer[K comparable](from []K, to []K) *sliceComparer[K] { + return &sliceComparer[K]{ + from: from, + to: to, + } +} + +// added returns the items which are present in `to` but not in `from`. +func (s *sliceComparer[K]) added() []K { + return thingsAdded(s.from, s.to) +} + +// removed returns the items which are present in `from` but not in `to`. +func (s *sliceComparer[K]) removed() []K { + return thingsAdded(s.to, s.from) +} + +// same returns the items which are in both `to` and `from`. +func (s *sliceComparer[K]) same() []K { + var same []K + for _, fromThing := range s.from { + for _, toThing := range s.to { + if fromThing == toThing { + same = append(same, fromThing) + break + } + } + } + return same +} + +// thingsAdded returns the comparable things which are present in `to` but not +// in `from`. +func thingsAdded[K comparable](from []K, to []K) []K { + var added []K + for i := range to { + var found bool + for j := range from { + if from[j] == to[i] { + found = true + break + } + } + if !found { + added = append(added, to[i]) + } + } + return added +} + +// partitionsComparer is used to compare the differences between two maps of +// table:[]partition. +type partitionsComparer struct { + from map[dax.TableKey]dax.Partitions + to map[dax.TableKey]dax.Partitions +} + +func newPartitionsComparer(from map[dax.TableKey]dax.Partitions, to map[dax.TableKey]dax.Partitions) *partitionsComparer { + return &partitionsComparer{ + from: from, + to: to, + } +} + +// added returns the partitions which are present in `to` but not in `from`. The +// results remain in the format of a map of table:[]partition. +func (p *partitionsComparer) added() map[dax.TableKey]dax.Partitions { + return partitionsAdded(p.from, p.to) +} + +// removed returns the partitions which are present in `from` but not in `to`. +// The results remain in the format of a map of table:[]partition. +func (p *partitionsComparer) removed() map[dax.TableKey]dax.Partitions { + return partitionsAdded(p.to, p.from) +} + +// partitionsAdded returns the partitions which are present in `to` but not in `from`. +func partitionsAdded(from map[dax.TableKey]dax.Partitions, to map[dax.TableKey]dax.Partitions) map[dax.TableKey]dax.Partitions { + if from == nil { + return to + } + + added := make(map[dax.TableKey]dax.Partitions) + for tt, tps := range to { + fps, found := from[tt] + if !found { + added[tt] = tps + continue + } + + addedPartitions := dax.Partitions{} + for i := range tps { + var found bool + for j := range fps { + if fps[j] == tps[i] { + found = true + break + } + } + if !found { + addedPartitions = append(addedPartitions, tps[i]) + } + } + + if len(addedPartitions) > 0 { + added[tt] = addedPartitions + } + } + return added +} + +// fieldsComparer is used to compare the differences between two maps of +// table:[]fieldVersion. +type fieldsComparer struct { + from map[dax.TableKey]dax.FieldVersions + to map[dax.TableKey]dax.FieldVersions +} + +func newFieldsComparer(from map[dax.TableKey]dax.FieldVersions, to map[dax.TableKey]dax.FieldVersions) *fieldsComparer { + return &fieldsComparer{ + from: from, + to: to, + } +} + +// added returns the fields which are present in `to` but not in `from`. The +// results remain in the format of a map of table:[]field. +func (f *fieldsComparer) added() map[dax.TableKey]dax.FieldVersions { + return fieldsAdded(f.from, f.to) +} + +// removed returns the fields which are present in `from` but not in `to`. +// The results remain in the format of a map of table:[]field. +func (f *fieldsComparer) removed() map[dax.TableKey]dax.FieldVersions { + return fieldsAdded(f.to, f.from) +} + +// fieldsAdded returns the fields which are present in `to` but not in `from`. +func fieldsAdded(from map[dax.TableKey]dax.FieldVersions, to map[dax.TableKey]dax.FieldVersions) map[dax.TableKey]dax.FieldVersions { + if from == nil { + return to + } + + added := make(map[dax.TableKey]dax.FieldVersions) + for tt, tps := range to { + fps, found := from[tt] + if !found { + added[tt] = tps + continue + } + + addedFieldVersions := dax.FieldVersions{} + for i := range tps { + var found bool + for j := range fps { + if fps[j] == tps[i] { + found = true + break + } + } + if !found { + addedFieldVersions = append(addedFieldVersions, tps[i]) + } + } + + if len(addedFieldVersions) > 0 { + added[tt] = addedFieldVersions + } + } + return added +} + +// shardsComparer is used to compare the differences between two maps of +// table:[]shardV. +type shardsComparer struct { + from map[dax.TableKey]dax.Shards + to map[dax.TableKey]dax.Shards +} + +func newShardsComparer(from map[dax.TableKey]dax.Shards, to map[dax.TableKey]dax.Shards) *shardsComparer { + return &shardsComparer{ + from: from, + to: to, + } +} + +// added returns the shards which are present in `to` but not in `from`. The +// results remain in the format of a map of table:[]shard. +func (s *shardsComparer) added() map[dax.TableKey]dax.Shards { + return shardsAdded(s.from, s.to) +} + +// removed returns the shards which are present in `from` but not in `to`. The +// results remain in the format of a map of table:[]shard. +func (s *shardsComparer) removed() map[dax.TableKey]dax.Shards { + return shardsAdded(s.to, s.from) +} + +// shardsAdded returns the shards which are present in `to` but not in `from`. +func shardsAdded(from map[dax.TableKey]dax.Shards, to map[dax.TableKey]dax.Shards) map[dax.TableKey]dax.Shards { + if from == nil { + return to + } + + added := make(map[dax.TableKey]dax.Shards) + for tt, tss := range to { + fss, found := from[tt] + if !found { + added[tt] = tss + continue + } + + addedShards := dax.Shards{} + for i := range tss { + var found bool + for j := range fss { + if fss[j] == tss[i] { + found = true + break + } + } + if !found { + addedShards = append(addedShards, tss[i]) + } + } + + if len(addedShards) > 0 { + added[tt] = addedShards + } + } + return added +} + +// createTableAndFields creates the FeatureBase Tables and Fields provided in +// the dax.Directive format. +func (api *API) createTableAndFields(tbl *dax.QualifiedTable, partitions dax.Partitions) error { + cim := &CreateIndexMessage{ + Index: string(tbl.Key()), + CreatedAt: 0, + Meta: IndexOptions{ + Keys: tbl.StringKeys(), + TrackExistence: true, + }, + } + + // Create the index in etcd as the system of record. + if err := api.holder.persistIndex(context.Background(), cim); err != nil { + return errors.Wrap(err, "persisting index") + } + + idx, err := api.holder.createIndexWithPartitions(cim, partitions) + if err != nil { + return errors.Wrapf(err, "adding index: %s", tbl.Name) + } + + // Add the fields + for _, fld := range tbl.Fields { + if fld.IsPrimaryKey() { + continue + } + if err := createField(idx, fld); err != nil { + return errors.Wrapf(err, "creating field: %s", fld.Name) + } + } + + return nil +} + +// createField creates a FeatureBase Field in the provided FeatureBase Index +// based on the provided field's type. +// +// TODO: `time` fields +func createField(idx *Index, fld *dax.Field) error { + // Set the cache type and size (or use default) for those fields which + // require them. + cacheType := DefaultCacheType + cacheSize := uint32(DefaultCacheSize) + if fld.Options.CacheType != "" { + cacheType = fld.Options.CacheType + cacheSize = fld.Options.CacheSize + } + + opts := []FieldOption{} + + switch fld.Type { + case dax.FieldTypeBool: + opts = append(opts, + OptFieldTypeBool(), + ) + case dax.FieldTypeDecimal: + opts = append(opts, + OptFieldTypeDecimal(fld.Options.Scale), + ) + case dax.FieldTypeID: + opts = append(opts, + OptFieldTypeMutex(cacheType, cacheSize), + ) + case dax.FieldTypeIDSet: + opts = append(opts, + OptFieldTypeSet(cacheType, cacheSize), + ) + case dax.FieldTypeInt: + opts = append(opts, + OptFieldTypeInt(fld.Options.Min.ToInt64(0), fld.Options.Max.ToInt64(0)), + ) + case dax.FieldTypeString: + opts = append(opts, + OptFieldTypeMutex(cacheType, cacheSize), + OptFieldKeys(), + ) + case dax.FieldTypeStringSet: + opts = append(opts, + OptFieldTypeSet(cacheType, cacheSize), + OptFieldKeys(), + ) + case dax.FieldTypeTimestamp: + opts = append(opts, + OptFieldTypeTimestamp(fld.Options.Epoch, fld.Options.TimeUnit), + ) + default: + return errors.Errorf("unsupport field type: %s", fld.Type) + } + + if _, err := idx.CreateField(string(fld.Name), "", opts...); err != nil { + return errors.Wrapf(err, "creating field on index: %s", fld.Name) + } + return nil +} diff --git a/api_directive_internal_test.go b/api_directive_internal_test.go new file mode 100644 index 000000000..12a3b8917 --- /dev/null +++ b/api_directive_internal_test.go @@ -0,0 +1,25 @@ +package pilosa + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestThingsAddedGeneric(t *testing.T) { + from := []string{"a", "b", "c"} + to := []string{"b", "c", "d"} + + added := thingsAdded(from, to) + assert.Equal(t, added, []string{"d"}) +} + +func TestSliceComparer(t *testing.T) { + from := []string{"a", "b", "c"} + to := []string{"b", "c", "d"} + + sc := newSliceComparer(from, to) + + added := sc.added() + assert.Equal(t, added, []string{"d"}) +} diff --git a/api_directive_test.go b/api_directive_test.go new file mode 100644 index 000000000..8d5b732e1 --- /dev/null +++ b/api_directive_test.go @@ -0,0 +1,97 @@ +package pilosa_test + +import ( + "context" + "testing" + + pilosa "github.com/molecula/featurebase/v3" + "github.com/molecula/featurebase/v3/dax" + daxtest "github.com/molecula/featurebase/v3/dax/test" + "github.com/molecula/featurebase/v3/test" + "github.com/stretchr/testify/assert" +) + +// Ensure holder can handle an incoming directive. +func TestAPI_Directive(t *testing.T) { + + c := test.MustRunCluster(t, 1) + defer c.Close() + + api := c.GetPrimary().API + ctx := context.Background() + + qual := dax.NewTableQualifier("acme", "db1") + tbl1 := daxtest.TestQualifiedTableWithID(t, qual, "1", "tbl1", 12, false) + tbl2 := daxtest.TestQualifiedTableWithID(t, qual, "2", "tbl2", 12, false) + tbl3 := daxtest.TestQualifiedTableWithID(t, qual, "3", "tbl3", 12, false) + + t.Run("Schema", func(t *testing.T) { + + // Empty directive (and empty holder). + { + d := &dax.Directive{ + Method: dax.DirectiveMethodDiff, + Version: 1, + } + err := api.ApplyDirective(ctx, d) + assert.NoError(t, err) + assertTablesMatch(t, []string{}, api.Holder().Indexes()) + } + + // Add a new table. + { + d := &dax.Directive{ + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbl1, + }, + Version: 2, + } + err := api.ApplyDirective(ctx, d) + assert.NoError(t, err) + assertTablesMatch(t, []string{"tbl__acme__db1__1"}, api.Holder().Indexes()) + } + + // Add a new table, and keep the existing table. + { + d := &dax.Directive{ + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbl1, + tbl2, + }, + Version: 3, + } + err := api.ApplyDirective(ctx, d) + assert.NoError(t, err) + assertTablesMatch(t, []string{"tbl__acme__db1__1", "tbl__acme__db1__2"}, api.Holder().Indexes()) + } + + // Add a new table and remove one of the existing tables. + { + d := &dax.Directive{ + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbl2, + tbl3, + }, + Version: 4, + } + err := api.ApplyDirective(ctx, d) + assert.NoError(t, err) + assertTablesMatch(t, []string{"tbl__acme__db1__2", "tbl__acme__db1__3"}, api.Holder().Indexes()) + } + }) +} + +// assertTablesMatch is a helper function which asserts that the list of index +// names in `actual` match those provided in `expected`. +func assertTablesMatch(t *testing.T, expected []string, actual []*pilosa.Index) { + t.Helper() + + act := make([]string, len(actual)) + for i := range actual { + act[i] = actual[i].Name() + } + assert.ElementsMatch(t, expected, act) +} diff --git a/batch/Makefile b/batch/Makefile index dac750101..0329794d3 100644 --- a/batch/Makefile +++ b/batch/Makefile @@ -7,11 +7,6 @@ GO ?= go # we might be running multiple instances of the tests concurrently. PROJECT ?= batch DOCKER_COMPOSE = docker-compose -p $(PROJECT) -BRANCH_NAME ?= "" - -.pulled: - $(DOCKER_COMPOSE) pull - touch .pulled vendor: ../go.mod $(GO) mod vendor @@ -19,24 +14,19 @@ vendor: ../go.mod build-%: $(DOCKER_COMPOSE) build $* -pull-%: - $(DOCKER_COMPOSE) pull $* - test-all: $(MAKE) startup $(MAKE) test-run $(MAKE) shutdown -start-all: .pulled build-wait - echo "branch name" ${BRANCH_NAME} - BRANCH_NAME=${BRANCH_NAME} $(DOCKER_COMPOSE) up -d featurebase +start-all: build-wait + $(DOCKER_COMPOSE) up -d featurebase $(DOCKER_COMPOSE) run -T wait featurebase curl --silent --fail http://featurebase:10101/status startup: start-all shutdown: $(DOCKER_COMPOSE) down -v --remove-orphans - rm -f .pulled save-%-logs: $(DOCKER_COMPOSE) logs $* > ./testdata/$(PROJECT)_$*_logs.txt @@ -48,7 +38,7 @@ test-run-local: $(DOCKER_COMPOSE) build batch-test $(DOCKER_COMPOSE) run -T batch-test go test -mod=vendor -tags=odbc,dynamic $(TCMD) -TPKG ?= ./... +TPKG ?= ../... test-run: vendor $(DOCKER_COMPOSE) build batch-test - $(DOCKER_COMPOSE) run -T batch-test bash -c "set -o pipefail; go test -v -mod=vendor -tags=odbc,dynamic $(TPKG) -covermode=atomic -coverpkg=$(TPKG) -json -coverprofile=/testdata/$(PROJECT)_base_coverage.out | tee /testdata/$(PROJECT)_report.out" + $(DOCKER_COMPOSE) run -T batch-test bash -c "set -o pipefail; go test -v -mod=vendor -tags=odbc,dynamic ./... -covermode=atomic -coverpkg=$(TPKG) -json -coverprofile=/testdata/$(PROJECT)_base_coverage.out | tee /testdata/$(PROJECT)_report.out" diff --git a/batch/README.md b/batch/README.md index 4ef894a83..fd67ad770 100644 --- a/batch/README.md +++ b/batch/README.md @@ -22,7 +22,7 @@ In addition to these dependancies, you will need to be added to the molecula [Gi First start the test environment. This is a docker-compose environment that includes featurebase. - BRANCH_NAME=master make startup + make startup To build and run the integration tests, run: diff --git a/batch/batch.go b/batch/batch.go index 4f01a1857..ad253d83b 100644 --- a/batch/batch.go +++ b/batch/batch.go @@ -1,3 +1,4 @@ +// Copyright 2021 Molecula Corp. All rights reserved. // Package batch provides tooling to prepare batches of records for ingest. package batch @@ -182,6 +183,8 @@ type Batch struct { clearFrags fragments useShardTransactionalEndpoint bool + + mdsHost string } func (b *Batch) Len() int { return len(b.ids) } @@ -332,6 +335,7 @@ func NewBatch(importer Importer, size int, index *featurebase.IndexInfo, fields return nil, errors.Wrap(err, "applying options") } } + return b, nil } diff --git a/batch/batch_test.go b/batch/batch_test.go index 442f8aa38..c5d33f5e8 100644 --- a/batch/batch_test.go +++ b/batch/batch_test.go @@ -1,3 +1,5 @@ +// Copyright 2021 Molecula Corp. All rights reserved. + package batch import ( @@ -1842,7 +1844,7 @@ func testImportBatchSetsAndClears(t *testing.T, importer Importer, sapi featureb // testTopNCacheRegression recreates an issue we saw in an IDK test // where if a value is completely removed (all bits unset from a row), -// it didn't get removed from the cache beacuse a full recalculation +// it didn't get removed from the cache because a full recalculation // had no way to clear the cache, it would just reset existing // values. We added Clear on the cache interface to fix this. func testTopNCacheRegression(t *testing.T, importer Importer, sapi featurebase.SchemaAPI, qapi featurebase.QueryAPI) { diff --git a/batch/docker-compose.yml b/batch/docker-compose.yml index 09b2c108c..b3a897c04 100644 --- a/batch/docker-compose.yml +++ b/batch/docker-compose.yml @@ -4,12 +4,13 @@ services: featurebase: build: context: ../. - dockerfile: ./Dockerfile + dockerfile: ./Dockerfile-clustertests environment: PILOSA_DATA_DIR: /data PILOSA_BIND: 0.0.0.0:10101 PILOSA_BIND_GRPC: 0.0.0.0:20101 PILOSA_ADVERTISE: featurebase:10101 + command: /featurebase -test.run=TestRunMain -test.coverprofile=/testdata/batch_coverage.out server volumes: - ./testdata:/testdata @@ -21,6 +22,8 @@ services: - ./testdata:/testdata wait: + depends_on: + - "featurebase" build: context: . dockerfile: Dockerfile-wait diff --git a/bsi.go b/bsi.go index 3fee24de2..c74958d8c 100644 --- a/bsi.go +++ b/bsi.go @@ -8,14 +8,14 @@ import ( "github.com/featurebasedb/featurebase/v3/roaring" ) -// bsiData contains BSI-structured data. -type bsiData []*Row +// BSIData contains BSI-structured data. +type BSIData []*Row -// pivotDescending loops over nonzero BSI values in descending order. +// PivotDescending loops over nonzero BSI values in descending order. // For each value, the provided function is called with the value and a slice of the associated columns. // If limit or offset are not-nil, they will be applied. // Applying a limit or offset may modify the pointed-to value. -func (bsi bsiData) pivotDescending(filter *Row, branch uint64, limit, offset *uint64, fn func(uint64, ...uint64)) { +func (bsi BSIData) PivotDescending(filter *Row, branch uint64, limit, offset *uint64, fn func(uint64, ...uint64)) { // This "pivot" algorithm works by treating the BSI data as a tree. // Each branch of this tree corresponds to a power-of-2-sized range of BSI values. // Each range is subdivided into 2 ranges of half size, which form lower branches. @@ -56,8 +56,8 @@ func (bsi bsiData) pivotDescending(filter *Row, branch uint64, limit, offset *ui upperBranch, lowerBranch := branch|(1< 0 { if _, err := w.Write([]byte("\n")); err != nil { return errors.Wrapf(err, "writing warning: %s", r.Error) @@ -228,12 +416,15 @@ func (r *response) WriteWarnings(w io.Writer) error { return nil } -func (r *response) WriteOut(w io.Writer) error { +func WriteOut(r *featurebase.SQLResponse, w io.Writer) error { + if r == nil { + return errors.New("attempt to write out nil response") + } if r.Error != "" { if _, err := w.Write([]byte("Error: " + r.Error + "\n")); err != nil { return errors.Wrapf(err, "writing error: %s", r.Error) } - return r.WriteWarnings(w) + return WriteWarnings(r, w) } t := table.NewWriter() @@ -255,7 +446,7 @@ func (r *response) WriteOut(w io.Writer) error { } t.Render() - err := r.WriteWarnings(w) + err := WriteWarnings(r, w) if err != nil { return err } @@ -282,3 +473,77 @@ func schemaToRow(schema featurebase.SQLSchema) []interface{} { } return ret } + +// Ensure type implements interface. +var _ FBQueryer = (*standardQueryer)(nil) + +// standardQueryer supports a standard featurebase deployment hitting the /sql +// endpoint with a payload containing only the sql statement. +type standardQueryer struct { + Host string + Port string +} + +func (qryr *standardQueryer) Query(org, db, sql string) (*featurebase.SQLResponse, error) { + buf := bytes.Buffer{} + url := fmt.Sprintf("%s/sql", hostPort(qryr.Host, qryr.Port)) + + buf.Write([]byte(sql)) + + resp, err := http.Post(url, "application/json", &buf) + if err != nil { + return nil, errors.Wrapf(err, "posting query") + } + + fullbod, err := io.ReadAll(resp.Body) + if err != nil { + return nil, errors.Wrap(err, "reading response") + } + sqlResponse := &featurebase.SQLResponse{} + if err := json.Unmarshal(fullbod, sqlResponse); err != nil { + return nil, errors.Wrapf(err, "unmarshaling query response, body:\n'%s'\n", fullbod) + } + + return sqlResponse, nil +} + +// Ensure type implements interface. +var _ FBQueryer = (*daxQueryer)(nil) + +// daxQueryer is similar to the standardQueryer except that it hits a different +// endpoint, and its payload is a json object which includes, in addition to the +// sql statement, things like org and db. +type daxQueryer struct { + Host string + Port string +} + +func (qryr *daxQueryer) Query(org, db, sql string) (*featurebase.SQLResponse, error) { + buf := bytes.Buffer{} + url := fmt.Sprintf("%s/queryer/sql", hostPort(qryr.Host, qryr.Port)) + + sqlReq := &queryerhttp.SQLRequest{ + OrganizationID: dax.OrganizationID(org), + DatabaseID: dax.DatabaseID(db), + SQL: sql, + } + if err := json.NewEncoder(&buf).Encode(sqlReq); err != nil { + return nil, errors.Wrapf(err, "encoding sql request: %s", sql) + } + + resp, err := http.Post(url, "application/json", &buf) + if err != nil { + return nil, errors.Wrapf(err, "posting query") + } + + fullbod, err := io.ReadAll(resp.Body) + if err != nil { + return nil, errors.Wrap(err, "reading response") + } + sqlResponse := &featurebase.SQLResponse{} + if err := json.Unmarshal(fullbod, sqlResponse); err != nil { + return nil, errors.Wrapf(err, "unmarshaling query response, body:\n'%s'\n", fullbod) + } + + return sqlResponse, nil +} diff --git a/ctl/dax.go b/ctl/dax.go new file mode 100644 index 000000000..a1d2f02c9 --- /dev/null +++ b/ctl/dax.go @@ -0,0 +1,39 @@ +package ctl + +import ( + "github.com/molecula/featurebase/v3/dax/server" + "github.com/spf13/cobra" +) + +// BuildDAXFlags attaches a set of flags to the command for a server instance. +func BuildDAXFlags(cmd *cobra.Command, srv *server.Command) { + flags := cmd.Flags() + + flags.StringVarP(&srv.Config.Bind, "bind", "b", srv.Config.Bind, "Default URI on which this service should listen.") + flags.StringVar(&srv.Config.Advertise, "advertise", srv.Config.Advertise, "Address to advertise externally.") + flags.BoolVar(&srv.Config.Verbose, "verbose", srv.Config.Verbose, "Enable verbose logging") + flags.StringVar(&srv.Config.LogPath, "log-path", srv.Config.LogPath, "Log path") + + flags.StringVar(&srv.Config.StorageMethod, "storage-method", srv.Config.StorageMethod, "Method to use for persistent storage.") + flags.StringVar(&srv.Config.StorageDSN, "storage-dsn", srv.Config.StorageDSN, "Datasource Name when using an applicable storage method.") + + // MDS + flags.BoolVar(&srv.Config.MDS.Run, "mds.run", srv.Config.MDS.Run, "Run the MDS service in process.") + flags.DurationVar(&srv.Config.MDS.Config.RegistrationBatchTimeout, "mds.config.registration-batch-timeout", srv.Config.MDS.Config.RegistrationBatchTimeout, "Timeout for node registration batches.") + + // WriteLogger + flags.BoolVar(&srv.Config.WriteLogger.Run, "writelogger.run", srv.Config.WriteLogger.Run, "Run the WriteLogger service in process.") + flags.StringVar(&srv.Config.WriteLogger.Config.DataDir, "writelogger.config.data-dir", srv.Config.WriteLogger.Config.DataDir, "WriteLogger directory to use in process.") + + // Snapshotter + flags.BoolVar(&srv.Config.Snapshotter.Run, "snapshotter.run", srv.Config.Snapshotter.Run, "Run the Snapshotter service in process.") + flags.StringVar(&srv.Config.Snapshotter.Config.DataDir, "snapshotter.config.data-dir", srv.Config.Snapshotter.Config.DataDir, "Snapshotter directory to use in process.") + + // Queryer + flags.BoolVar(&srv.Config.Queryer.Run, "queryer.run", srv.Config.Queryer.Run, "Run the Queryer service in process.") + flags.StringVar(&srv.Config.Queryer.Config.MDSAddress, "queryer.config.mds-address", srv.Config.Queryer.Config.MDSAddress, "Address of remote MDS process.") + + // Computer + flags.BoolVar(&srv.Config.Computer.Run, "computer.run", srv.Config.Computer.Run, "Run the Computer service in process.") + flags.AddFlagSet(serverFlagSet(&srv.Config.Computer.Config, "computer.config")) +} diff --git a/ctl/server.go b/ctl/server.go index ea8cc0a89..8463dcf28 100644 --- a/ctl/server.go +++ b/ctl/server.go @@ -8,114 +8,146 @@ import ( "github.com/featurebasedb/featurebase/v3/server" "github.com/featurebasedb/featurebase/v3/storage" "github.com/spf13/cobra" + "github.com/spf13/pflag" ) +// serverFlagSet returns a pflag.FlagSet. All flags will be prefixed with the +// given prefix value, and default values come from the provided server.Config. +func serverFlagSet(srv *server.Config, prefix string) *pflag.FlagSet { + // pre applies prefix to s when a prefix is provided. + pre := func(s string) string { + if prefix == "" { + return s + } + return prefix + "." + s + } + + // short will pass through the short flag as long as a prefix is not + // specified. + short := func(s string) string { + if prefix == "" { + return s + } + return "" + } + + flags := pflag.NewFlagSet("featurebase", pflag.ExitOnError) + flags.StringVar(&srv.Name, pre("name"), srv.Name, "Name of the node in the cluster.") + flags.StringVar(&srv.MDSAddress, pre("mds-address"), srv.MDSAddress, "MDS service to register with.") + flags.StringVar(&srv.WriteLogger, pre("write-logger"), srv.WriteLogger, "WriteLogger to read/write append logs.") + flags.StringVar(&srv.Snapshotter, pre("snapshotter"), srv.Snapshotter, "Snapshotter to read/write snapshots.") + flags.StringVarP(&srv.DataDir, pre("data-dir"), short("d"), srv.DataDir, "Directory to store FeatureBase data files.") + flags.StringVarP(&srv.Bind, pre("bind"), short("b"), srv.Bind, "Default URI on which FeatureBase should listen.") + flags.StringVar(&srv.BindGRPC, pre("bind-grpc"), srv.BindGRPC, "URI on which FeatureBase should listen for gRPC requests.") + flags.StringVar(&srv.Advertise, pre("advertise"), srv.Advertise, "Address to advertise externally.") + flags.StringVar(&srv.AdvertiseGRPC, pre("advertise-grpc"), srv.AdvertiseGRPC, "Address to advertise externally for gRPC.") + flags.IntVar(&srv.MaxWritesPerRequest, pre("max-writes-per-request"), srv.MaxWritesPerRequest, "Number of write commands per request.") + flags.StringVar(&srv.LogPath, pre("log-path"), srv.LogPath, "Log path") + flags.BoolVar(&srv.Verbose, pre("verbose"), srv.Verbose, "Enable verbose logging") + flags.Uint64Var(&srv.MaxMapCount, pre("max-map-count"), srv.MaxMapCount, "Limits the maximum number of active mmaps. FeatureBase will fall back to reading files once this is exhausted. Set below your system's vm.max_map_count.") + flags.Uint64Var(&srv.MaxFileCount, pre("max-file-count"), srv.MaxFileCount, "Soft limit on the maximum number of fragment files FeatureBase keeps open simultaneously.") + flags.DurationVar((*time.Duration)(&srv.LongQueryTime), pre("long-query-time"), time.Duration(srv.LongQueryTime), "Duration that will trigger log and stat messages for slow queries. Zero to disable.") + flags.IntVar(&srv.QueryHistoryLength, pre("query-history-length"), srv.QueryHistoryLength, "Number of queries to remember in history.") + flags.Int64Var(&srv.MaxQueryMemory, pre("max-query-memory"), srv.MaxQueryMemory, "Maximum memory allowed per Extract() or SELECT query.") + + // TLS + SetTLSConfig(flags, pre(""), &srv.TLS.CertificatePath, &srv.TLS.CertificateKeyPath, &srv.TLS.CACertPath, &srv.TLS.SkipVerify, &srv.TLS.EnableClientVerification) + + // Handler + flags.StringSliceVar(&srv.Handler.AllowedOrigins, pre("handler.allowed-origins"), []string{}, "Comma separated list of allowed origin URIs (for CORS/Web UI).") + + // Cluster + flags.IntVar(&srv.Cluster.ReplicaN, pre("cluster.replicas"), 1, "Number of hosts each piece of data should be stored on.") + flags.DurationVar((*time.Duration)(&srv.Cluster.LongQueryTime), pre("cluster.long-query-time"), time.Duration(srv.Cluster.LongQueryTime), "RENAMED TO 'long-query-time': Duration that will trigger log and stat messages for slow queries.") // negative duration indicates invalid value because 0 is meaningful + flags.StringVar(&srv.Cluster.Name, pre("cluster.name"), srv.Cluster.Name, "Human-readable name for the cluster.") + flags.StringVar(&srv.Cluster.PartitionToNodeAssignment, pre("cluster.partition-to-node-assignment"), srv.Cluster.PartitionToNodeAssignment, "How to assign partitions to nodes. jmp-hash or modulus") + + // Translation + flags.StringVar(&srv.Translation.PrimaryURL, pre("translation.primary-url"), srv.Translation.PrimaryURL, "DEPRECATED: URL for primary translation node for replication.") + flags.IntVar(&srv.Translation.MapSize, pre("translation.map-size"), srv.Translation.MapSize, "Size in bytes of mmap to allocate for key translation.") + + // Etcd + // Etcd.Name used Config.Name for its value. + flags.StringVar(&srv.Etcd.Dir, pre("etcd.dir"), srv.Etcd.Dir, "Directory to store etcd data files. If not provided, a directory will be created under the main data-dir directory.") + // Etcd.ClusterName uses Cluster.Name for its value + flags.StringVar(&srv.Etcd.LClientURL, pre("etcd.listen-client-address"), srv.Etcd.LClientURL, "Listen client address.") + flags.StringVar(&srv.Etcd.AClientURL, pre("etcd.advertise-client-address"), srv.Etcd.AClientURL, "Advertise client address. If not provided, uses the listen client address.") + flags.StringVar(&srv.Etcd.LPeerURL, pre("etcd.listen-peer-address"), srv.Etcd.LPeerURL, "Listen peer address.") + flags.StringVar(&srv.Etcd.APeerURL, pre("etcd.advertise-peer-address"), srv.Etcd.APeerURL, "Advertise peer address. If not provided, uses the listen peer address.") + flags.StringVar(&srv.Etcd.ClusterURL, pre("etcd.cluster-url"), srv.Etcd.ClusterURL, "Cluster URL to join.") + flags.StringVar(&srv.Etcd.InitCluster, pre("etcd.initial-cluster"), srv.Etcd.InitCluster, "Initial cluster name1=apurl1,name2=apurl2") + flags.Int64Var(&srv.Etcd.HeartbeatTTL, pre("etcd.heartbeat-ttl"), srv.Etcd.HeartbeatTTL, "Timeout used to determine cluster status") + + flags.StringVar(&srv.Etcd.Cluster, "etcd.static-cluster", srv.Etcd.Cluster, "EXPERIMENTAL static featurebase cluster name1=apurl1,name2=apurl2") + flags.MarkHidden("etcd.static-cluster") + flags.StringVar(&srv.Etcd.EtcdHosts, "etcd.etcd-hosts", srv.Etcd.EtcdHosts, "EXPERIMENTAL etcd server host:port comma separated list") + flags.MarkHidden("etcd.etcd-hosts") // TODO (twg) expose when ready for public consumption + + // External postgres database for ExternalLookup + flags.StringVar(&srv.LookupDBDSN, pre("lookup-db-dsn"), "", "external (postgres) database DSN to use for ExternalLookup calls") + + // AntiEntropy + flags.DurationVar((*time.Duration)(&srv.AntiEntropy.Interval), pre("anti-entropy.interval"), (time.Duration)(srv.AntiEntropy.Interval), "Interval at which to run anti-entropy routine.") + + // Metric + flags.StringVar(&srv.Metric.Service, pre("metric.service"), srv.Metric.Service, "Where to send stats: can be expvar (in-memory served at /debug/vars), prometheus, statsd or none.") + flags.StringVar(&srv.Metric.Host, pre("metric.host"), srv.Metric.Host, "URI to send metrics when metric.service is statsd.") + flags.DurationVar((*time.Duration)(&srv.Metric.PollInterval), pre("metric.poll-interval"), (time.Duration)(srv.Metric.PollInterval), "Polling interval metrics.") + flags.BoolVar((&srv.Metric.Diagnostics), pre("metric.diagnostics"), srv.Metric.Diagnostics, "Enabled diagnostics reporting.") + + // Tracing + flags.StringVar(&srv.Tracing.AgentHostPort, pre("tracing.agent-host-port"), srv.Tracing.AgentHostPort, "Jaeger agent host:port.") + flags.StringVar(&srv.Tracing.SamplerType, pre("tracing.sampler-type"), srv.Tracing.SamplerType, "Jaeger sampler type (remote, const, probabilistic, ratelimiting) or 'off' to disable tracing completely.") + flags.Float64Var(&srv.Tracing.SamplerParam, pre("tracing.sampler-param"), srv.Tracing.SamplerParam, "Jaeger sampler parameter.") + + // Profiling + flags.IntVar(&srv.Profile.BlockRate, pre("profile.block-rate"), srv.Profile.BlockRate, "Sampling rate for goroutine blocking profiler. One sample per ns.") + flags.IntVar(&srv.Profile.MutexFraction, pre("profile.mutex-fraction"), srv.Profile.MutexFraction, "Sampling fraction for mutex contention profiling. Sample 1/ of events.") + + flags.StringVar(&srv.Storage.Backend, pre("storage.backend"), storage.DefaultBackend, "Storage backend to use: 'rbf' is only supported value.") + flags.BoolVar(&srv.Storage.FsyncEnabled, pre("storage.fsync"), true, "enable fsync fully safe flush-to-disk") + + // RBF specific flags. See pilosa/rbf/cfg/cfg.go for definitions. + srv.RBFConfig.DefineFlags(flags, prefix) + + flags.BoolVar(&srv.SQL.EndpointEnabled, pre("sql.endpoint-enabled"), srv.SQL.EndpointEnabled, "Enable FeatureBase SQL /sql endpoint (default false)") + + flags.DurationVar(&srv.CheckInInterval, pre("check-in-interval"), srv.CheckInInterval, "Interval between check-ins to MDS") + + // Future flags. + flags.BoolVar(&srv.Future.Rename, pre("future.rename"), false, "Present application name as FeatureBase. Defaults to false, will default to true in an upcoming release.") + + // OAuth2.0 identity provider configuration + flags.BoolVar(&srv.Auth.Enable, pre("auth.enable"), false, "Enable AuthN/AuthZ of featurebase, disabled by default.") + flags.StringVar(&srv.Auth.ClientId, pre("auth.client-id"), srv.Auth.ClientId, "Identity Provider's Application/Client ID.") + flags.StringVar(&srv.Auth.ClientSecret, pre("auth.client-secret"), srv.Auth.ClientSecret, "Identity Provider's Client Secret.") + flags.StringVar(&srv.Auth.AuthorizeURL, pre("auth.authorize-url"), srv.Auth.AuthorizeURL, "Identity Provider's Authorize URL.") + flags.StringVar(&srv.Auth.RedirectBaseURL, pre("auth.redirect-base-url"), srv.Auth.RedirectBaseURL, "Base URL of the featurebase instance used to redirect IDP.") + flags.StringVar(&srv.Auth.TokenURL, pre("auth.token-url"), srv.Auth.TokenURL, "Identity Provider's Token URL.") + flags.StringVar(&srv.Auth.GroupEndpointURL, pre("auth.group-endpoint-url"), srv.Auth.GroupEndpointURL, "Identity Provider's Group endpoint URL.") + flags.StringVar(&srv.Auth.LogoutURL, pre("auth.logout-url"), srv.Auth.LogoutURL, "Identity Provider's Logout URL.") + flags.StringSliceVar(&srv.Auth.Scopes, pre("auth.scopes"), srv.Auth.Scopes, "Comma separated list of scopes obtained from IdP") + flags.StringVar(&srv.Auth.SecretKey, pre("auth.secret-key"), srv.Auth.SecretKey, "Secret key used for auth.") + flags.StringVar(&srv.Auth.PermissionsFile, pre("auth.permissions"), srv.Auth.PermissionsFile, "Permissions' file with group authorization.") + flags.StringVar(&srv.Auth.QueryLogPath, pre("auth.query-log-path"), srv.Auth.QueryLogPath, "Path to log user queries") + flags.StringSliceVar(&srv.Auth.ConfiguredIPs, pre("auth.configured-ips"), srv.Auth.ConfiguredIPs, "List of configured IPs allowed for ingest") + + flags.BoolVar(&srv.DataDog.Enable, pre("datadog.enable"), false, "enable continuous profiling with DataDog cloud service, Note you must have DataDog agent installed") + flags.StringVar(&srv.DataDog.Service, pre("datadog.service"), "default-service", "The Datadog service name, for example my-web-app") + flags.StringVar(&srv.DataDog.Env, pre("datadog.env"), "default-env", "The Datadog environment name, for example, production") + flags.StringVar(&srv.DataDog.Version, pre("datadog.version"), "default-version", "The version of your application") + flags.StringVar(&srv.DataDog.Tags, pre("datadog.tags"), "molecula", "The tags to apply to an uploaded profile. Must be a list of in the format :,:") + flags.BoolVar(&srv.DataDog.CPUProfile, pre("datadog.cpu-profile"), true, "golang pprof cpu profile ") + flags.BoolVar(&srv.DataDog.HeapProfile, pre("datadog.heap-profile"), true, "golang pprof heap profile") + flags.BoolVar(&srv.DataDog.MutexProfile, pre("datadog.mutex-profile"), false, "golang pprof mutex profile") + flags.BoolVar(&srv.DataDog.GoroutineProfile, pre("datadog.goroutine-profile"), false, "golang pprof goroutine profile") + flags.BoolVar(&srv.DataDog.BlockProfile, pre("datadog.block-profile"), false, "golang pprof goroutine ") + + return flags +} + // BuildServerFlags attaches a set of flags to the command for a server instance. func BuildServerFlags(cmd *cobra.Command, srv *server.Command) { flags := cmd.Flags() - flags.StringVar(&srv.Config.Name, "name", srv.Config.Name, "Name of the node in the cluster.") - flags.StringVarP(&srv.Config.DataDir, "data-dir", "d", srv.Config.DataDir, "Directory to store FeatureBase data files.") - flags.StringVarP(&srv.Config.Bind, "bind", "b", srv.Config.Bind, "Default URI on which FeatureBase should listen.") - flags.StringVar(&srv.Config.BindGRPC, "bind-grpc", srv.Config.BindGRPC, "URI on which FeatureBase should listen for gRPC requests.") - flags.StringVar(&srv.Config.Advertise, "advertise", srv.Config.Advertise, "Address to advertise externally.") - flags.StringVar(&srv.Config.AdvertiseGRPC, "advertise-grpc", srv.Config.AdvertiseGRPC, "Address to advertise externally for gRPC.") - flags.IntVar(&srv.Config.MaxWritesPerRequest, "max-writes-per-request", srv.Config.MaxWritesPerRequest, "Number of write commands per request.") - flags.StringVar(&srv.Config.LogPath, "log-path", srv.Config.LogPath, "Log path") - flags.BoolVar(&srv.Config.Verbose, "verbose", srv.Config.Verbose, "Enable verbose logging") - flags.Uint64Var(&srv.Config.MaxMapCount, "max-map-count", srv.Config.MaxMapCount, "Limits the maximum number of active mmaps. FeatureBase will fall back to reading files once this is exhausted. Set below your system's vm.max_map_count.") - flags.Uint64Var(&srv.Config.MaxFileCount, "max-file-count", srv.Config.MaxFileCount, "Soft limit on the maximum number of fragment files FeatureBase keeps open simultaneously.") - flags.DurationVar((*time.Duration)(&srv.Config.LongQueryTime), "long-query-time", time.Duration(srv.Config.LongQueryTime), "Duration that will trigger log and stat messages for slow queries. Zero to disable.") - flags.IntVar(&srv.Config.QueryHistoryLength, "query-history-length", srv.Config.QueryHistoryLength, "Number of queries to remember in history.") - flags.Int64Var(&srv.Config.MaxQueryMemory, "max-query-memory", srv.Config.MaxQueryMemory, "Maximum memory allowed per Extract() or SELECT query.") - - // TLS - SetTLSConfig(flags, "", &srv.Config.TLS.CertificatePath, &srv.Config.TLS.CertificateKeyPath, &srv.Config.TLS.CACertPath, &srv.Config.TLS.SkipVerify, &srv.Config.TLS.EnableClientVerification) - - // Handler - flags.StringSliceVar(&srv.Config.Handler.AllowedOrigins, "handler.allowed-origins", []string{}, "Comma separated list of allowed origin URIs (for CORS/Web UI).") - - // Cluster - flags.IntVar(&srv.Config.Cluster.ReplicaN, "cluster.replicas", 1, "Number of hosts each piece of data should be stored on.") - flags.DurationVar((*time.Duration)(&srv.Config.Cluster.LongQueryTime), "cluster.long-query-time", time.Duration(srv.Config.Cluster.LongQueryTime), "RENAMED TO 'long-query-time': Duration that will trigger log and stat messages for slow queries.") // negative duration indicates invalid value because 0 is meaningful - flags.StringVar(&srv.Config.Cluster.Name, "cluster.name", srv.Config.Cluster.Name, "Human-readable name for the cluster.") - flags.StringVar(&srv.Config.Cluster.PartitionToNodeAssignment, "cluster.partition-to-node-assignment", srv.Config.Cluster.PartitionToNodeAssignment, "How to assign partitions to nodes. jmp-hash or modulus") - - // Translation - flags.StringVar(&srv.Config.Translation.PrimaryURL, "translation.primary-url", srv.Config.Translation.PrimaryURL, "DEPRECATED: URL for primary translation node for replication.") - flags.IntVar(&srv.Config.Translation.MapSize, "translation.map-size", srv.Config.Translation.MapSize, "Size in bytes of mmap to allocate for key translation.") - - // Etcd - // Etcd.Name used Config.Name for its value. - flags.StringVar(&srv.Config.Etcd.Dir, "etcd.dir", srv.Config.Etcd.Dir, "Directory to store etcd data files. If not provided, a directory will be created under the main data-dir directory.") - // Etcd.ClusterName uses Cluster.Name for its value - flags.StringVar(&srv.Config.Etcd.LClientURL, "etcd.listen-client-address", srv.Config.Etcd.LClientURL, "Listen client address.") - flags.StringVar(&srv.Config.Etcd.AClientURL, "etcd.advertise-client-address", srv.Config.Etcd.AClientURL, "Advertise client address. If not provided, uses the listen client address.") - flags.StringVar(&srv.Config.Etcd.LPeerURL, "etcd.listen-peer-address", srv.Config.Etcd.LPeerURL, "Listen peer address.") - flags.StringVar(&srv.Config.Etcd.APeerURL, "etcd.advertise-peer-address", srv.Config.Etcd.APeerURL, "Advertise peer address. If not provided, uses the listen peer address.") - flags.StringVar(&srv.Config.Etcd.ClusterURL, "etcd.cluster-url", srv.Config.Etcd.ClusterURL, "Cluster URL to join.") - flags.StringVar(&srv.Config.Etcd.InitCluster, "etcd.initial-cluster", srv.Config.Etcd.InitCluster, "Initial cluster name1=apurl1,name2=apurl2") - flags.Int64Var(&srv.Config.Etcd.HeartbeatTTL, "etcd.heartbeat-ttl", srv.Config.Etcd.HeartbeatTTL, "Timeout used to determine cluster status") - - flags.StringVar(&srv.Config.Etcd.Cluster, "etcd.static-cluster", srv.Config.Etcd.Cluster, "EXPERIMENTAL static featurebase cluster name1=apurl1,name2=apurl2") - _ = flags.MarkHidden("etcd.static-cluster") - flags.StringVar(&srv.Config.Etcd.EtcdHosts, "etcd.etcd-hosts", srv.Config.Etcd.EtcdHosts, "EXPERIMENTAL etcd server host:port comma separated list") - _ = flags.MarkHidden("etcd.etcd-hosts") // TODO (twg) expose when ready for public consumption - - // External postgres database for ExternalLookup - flags.StringVar(&srv.Config.LookupDBDSN, "lookup-db-dsn", "", "external (postgres) database DSN to use for ExternalLookup calls") - - // AntiEntropy - flags.DurationVar((*time.Duration)(&srv.Config.AntiEntropy.Interval), "anti-entropy.interval", (time.Duration)(srv.Config.AntiEntropy.Interval), "Interval at which to run anti-entropy routine.") - - // Metric - flags.StringVar(&srv.Config.Metric.Service, "metric.service", srv.Config.Metric.Service, "Where to send stats: can be expvar (in-memory served at /debug/vars), prometheus, statsd or none.") - flags.StringVar(&srv.Config.Metric.Host, "metric.host", srv.Config.Metric.Host, "URI to send metrics when metric.service is statsd.") - flags.DurationVar((*time.Duration)(&srv.Config.Metric.PollInterval), "metric.poll-interval", (time.Duration)(srv.Config.Metric.PollInterval), "Polling interval metrics.") - flags.BoolVar((&srv.Config.Metric.Diagnostics), "metric.diagnostics", srv.Config.Metric.Diagnostics, "Enabled diagnostics reporting.") - - // Tracing - flags.StringVar(&srv.Config.Tracing.AgentHostPort, "tracing.agent-host-port", srv.Config.Tracing.AgentHostPort, "Jaeger agent host:port.") - flags.StringVar(&srv.Config.Tracing.SamplerType, "tracing.sampler-type", srv.Config.Tracing.SamplerType, "Jaeger sampler type (remote, const, probabilistic, ratelimiting) or 'off' to disable tracing completely.") - flags.Float64Var(&srv.Config.Tracing.SamplerParam, "tracing.sampler-param", srv.Config.Tracing.SamplerParam, "Jaeger sampler parameter.") - - // Profiling - flags.IntVar(&srv.Config.Profile.BlockRate, "profile.block-rate", srv.Config.Profile.BlockRate, "Sampling rate for goroutine blocking profiler. One sample per ns.") - flags.IntVar(&srv.Config.Profile.MutexFraction, "profile.mutex-fraction", srv.Config.Profile.MutexFraction, "Sampling fraction for mutex contention profiling. Sample 1/ of events.") - - flags.StringVar(&srv.Config.Storage.Backend, "storage.backend", storage.DefaultBackend, "Storage backend to use: 'rbf' is only supported value.") - flags.BoolVar(&srv.Config.Storage.FsyncEnabled, "storage.fsync", true, "enable fsync fully safe flush-to-disk") - - // RBF specific flags. See pilosa/rbf/cfg/cfg.go for definitions. - srv.Config.RBFConfig.DefineFlags(flags) - - flags.BoolVar(&srv.Config.SQL.EndpointEnabled, "sql.endpoint-enabled", srv.Config.SQL.EndpointEnabled, "Enable FeatureBase SQL /sql endpoint (default false)") - - // Future flags. - flags.BoolVar(&srv.Config.Future.Rename, "future.rename", false, "Present application name as FeatureBase. Defaults to false, will default to true in an upcoming release.") - - // OAuth2.0 identity provider configuration - flags.BoolVar(&srv.Config.Auth.Enable, "auth.enable", false, "Enable AuthN/AuthZ of featurebase, disabled by default.") - flags.StringVar(&srv.Config.Auth.ClientId, "auth.client-id", srv.Config.Auth.ClientId, "Identity Provider's Application/Client ID.") - flags.StringVar(&srv.Config.Auth.ClientSecret, "auth.client-secret", srv.Config.Auth.ClientSecret, "Identity Provider's Client Secret.") - flags.StringVar(&srv.Config.Auth.AuthorizeURL, "auth.authorize-url", srv.Config.Auth.AuthorizeURL, "Identity Provider's Authorize URL.") - flags.StringVar(&srv.Config.Auth.RedirectBaseURL, "auth.redirect-base-url", srv.Config.Auth.RedirectBaseURL, "Base URL of the featurebase instance used to redirect IDP.") - flags.StringVar(&srv.Config.Auth.TokenURL, "auth.token-url", srv.Config.Auth.TokenURL, "Identity Provider's Token URL.") - flags.StringVar(&srv.Config.Auth.GroupEndpointURL, "auth.group-endpoint-url", srv.Config.Auth.GroupEndpointURL, "Identity Provider's Group endpoint URL.") - flags.StringVar(&srv.Config.Auth.LogoutURL, "auth.logout-url", srv.Config.Auth.LogoutURL, "Identity Provider's Logout URL.") - flags.StringSliceVar(&srv.Config.Auth.Scopes, "auth.scopes", srv.Config.Auth.Scopes, "Comma separated list of scopes obtained from IdP") - flags.StringVar(&srv.Config.Auth.SecretKey, "auth.secret-key", srv.Config.Auth.SecretKey, "Secret key used for auth.") - flags.StringVar(&srv.Config.Auth.PermissionsFile, "auth.permissions", srv.Config.Auth.PermissionsFile, "Permissions' file with group authorization.") - flags.StringVar(&srv.Config.Auth.QueryLogPath, "auth.query-log-path", srv.Config.Auth.QueryLogPath, "Path to log user queries") - flags.StringSliceVar(&srv.Config.Auth.ConfiguredIPs, "auth.configured-ips", srv.Config.Auth.ConfiguredIPs, "List of configured IPs allowed for ingest") - - flags.BoolVar(&srv.Config.DataDog.Enable, "datadog.enable", false, "enable continuous profiling with DataDog cloud service, Note you must have DataDog agent installed") - flags.StringVar(&srv.Config.DataDog.Service, "datadog.service", "default-service", "The Datadog service name, for example my-web-app") - flags.StringVar(&srv.Config.DataDog.Env, "datadog.env", "default-env", "The Datadog environment name, for example, production") - flags.StringVar(&srv.Config.DataDog.Version, "datadog.version", "default-version", "The version of your application") - flags.StringVar(&srv.Config.DataDog.Tags, "datadog.tags", "molecula", "The tags to apply to an uploaded profile. Must be a list of in the format :,:") - flags.BoolVar(&srv.Config.DataDog.CPUProfile, "datadog.cpu-profile", true, "golang pprof cpu profile ") - flags.BoolVar(&srv.Config.DataDog.HeapProfile, "datadog.heap-profile", true, "golang pprof heap profile") - flags.BoolVar(&srv.Config.DataDog.MutexProfile, "datadog.mutex-profile", false, "golang pprof mutex profile") - flags.BoolVar(&srv.Config.DataDog.GoroutineProfile, "datadog.goroutine-profile", false, "golang pprof goroutine profile") - flags.BoolVar(&srv.Config.DataDog.BlockProfile, "datadog.block-profile", false, "golang pprof goroutine ") + flags.AddFlagSet(serverFlagSet(srv.Config, "")) } diff --git a/dax/Makefile b/dax/Makefile new file mode 100644 index 000000000..6eb0bc7f8 --- /dev/null +++ b/dax/Makefile @@ -0,0 +1,89 @@ +.PHONY: test testv test-integration testv-integration + +MCLOUD_ENV ?= sandbox +MCLOUD_ENV_FILE=.env.$(MCLOUD_ENV) +-include $(MCLOUD_ENV_FILE) + + +GO=go + +test: + $(GO) test ./... -short + +testv: + $(GO) test -v ./... -short + +test-integration: + mkdir -p ../coverage-from-docker + $(GO) test ./test/dax -count 1 -run Integration + +testv-integration: + $(GO) test -v ./test/dax -count 1 -run Integration + + + +############################### AWS STUFF ############################### + +AWS_REGION ?= +AWS_PROFILE ?= + +AWS = aws --profile=$(AWS_PROFILE) --region=$(AWS_REGION) + +# After pushing new images, use "make redeploy-ecs" to redeploy all DAX services. +redeploy-ecs: redeploy-svc-mds redeploy-svc-computer redeploy-svc-queryer + +redeploy-svc-%: + $(AWS) ecs update-service --cluster DAX --service $*-$(MCLOUD_ENV)-ecs-service --force-new-deployment --no-cli-pager + + +# Scale changes the desired count of the computer service. e.g. "make scale N=4" +scale: + $(AWS) ecs update-service --cluster DAX --service computer-$(MCLOUD_ENV)-ecs-service --desired-count=$(N) --no-cli-pager + + +I ?= 0 +# Get a shell on a running contianer. e.g. "make mds-shell", "make datagen-shell", etc. +%-shell: + $(eval TASK_ARN := $(shell $(AWS) ecs list-tasks --cluster=DAX --family=$*-family | jq -r .taskArns[$(I)])) + $(AWS) ecs execute-command --cluster=DAX --task=$(TASK_ARN) --command=/bin/sh --interactive + +datagen: task-arn-datagen + $(eval TASK_ARN := $(shell $(AWS) ecs list-tasks --cluster=DAX --family=$*-family | jq -r .taskArns[$(I)])) + $(AWS) ecs run-task --cluster=DAX --task-definition=$(TASK_ARN) --cli-input-json=file://./datagen_task_input.json --no-cli-pager --enable-execute-command + + +####################### docker-compose stuff ##############################3 + + +dc-reset: dc-prereqs dc-down + rm -rf ./dax-data/{snapshotter,writelogger}/* + +dc-build: + cd .. && $(MAKE) build-for-quick + docker-compose build + +dc-up: + docker-compose up -d + +dc-down: + docker-compose down + +dc-full-reup: dc-reset dc-build dc-up + +dc-logs: + docker-compose logs -f + +dc-logs-%: + docker-compose logs -f $* + +dc-prereqs: + mkdir -p ../.quick + +# This is just an example. For it to work, you'll first need to: +# featurebase cli --host localhost --port 8080 --org-id=testorg --db-id=testdb +# create table keysidstbl2 (_id string, slice idset); +dc-datagen: + docker-compose run datagen --end-at=500 --pilosa.batch-size=500 --featurebase.table-name=keysidstbl2 + +dc-exec-%: + docker-compose exec $* /bin/sh diff --git a/dax/README.md b/dax/README.md new file mode 100644 index 000000000..97c18b5e5 --- /dev/null +++ b/dax/README.md @@ -0,0 +1,64 @@ +# DAX + +DAX encapsulates anything which covers all of the services which make up the +"disaggregation of storage and compute" project. Initially, this will include +integration tests which pull in things like Metadata Services (MDS), including +the Controller, as well as FeatureBase and IDK-based ingesters. + +## Setting up the tests + +The DAX test currently requires docker images for: `featurebase` and `datagen`. + +If at any point you run into problems with go mod failing to reference a private +repo, make sure that you have `gitlab.com/molecula` in your `GOPRIVATE` +environment variable. + +Note that during the docker image build step, `go mod vendor` is run, which +creates a `vendor` directory in the root directory, and copies that to docker +during the build stage. Just be aware of this; you may want to remove that +vendor directory after you're done building docker images. + +#### Possible Configuruation +The following were relevent when the dax code was in a separate repository. +These may no longer be relevant. + +I needed to but this in my `~/.profile` file: + +```export GOPRIVATE=github.com/molecula,gitlab.com/molecula``` + +And this in my `~/.gitconfig` + +``` +[url "ssh://git@github.com/"] + insteadOf = https://github.com/ +[url "ssh://git@gitlab.com/"] + insteadOf = https://gitlab.com/ +``` + +Then `make docker` ran successfully. + + +### Build the FeatureBase docker image + +- Check out the + [dax](https://github.com/molecula/featurebase/tree/dax) + branch of the + [featurebase](https://github.com/molecula/featurebase) repository. +- Run `make docker-image-featurebase` to build the docker image +- You should now have an image in docker named `dax/featurebase` with the tag `latest`. + +### Build the Datagen docker image + +- `cd /idk` +- Run `make docker-image-datagen` to build the docker image +- You should now have an image in docker named `dax/datagen` with the tag + `latest`. + +## Running the tests + +- Check out the + [dax](https://github.com/molecula/featurebase/tree/dax) + branch of the + [featurebase](https://github.com/molecula/featurebase) repository. +- Change into the `dax` directory: `cd dax` +- Run `make test-integration`. diff --git a/dax/address.go b/dax/address.go new file mode 100644 index 000000000..b63f5f950 --- /dev/null +++ b/dax/address.go @@ -0,0 +1,138 @@ +package dax + +import ( + "context" + "fmt" + "strconv" + "strings" +) + +// Address is a string of the form [scheme]://[host]:[port] +type Address string + +// String returns the Address as a string type. +func (a Address) String() string { + return string(a) +} + +// Scheme returns the [scheme] portion of the Address. This may be an empty +// string if Address does not contain a scheme. +func (a Address) Scheme() string { + return parse(a).scheme +} + +// HostPort returns the [host]:[port] portion of the Address; in other words, +// the Address stripped of any scheme. +func (a Address) HostPort() string { + return parse(a).hostPort() +} + +// Host returns the [host] portion of the Address. +func (a Address) Host() string { + return parse(a).host +} + +// Port returns the [port] portion of the Address. If the port values is invalid +// or does not exist, the returned value will default to 0. +func (a Address) Port() uint16 { + return parse(a).port +} + +// OverrideScheme overrides Address's current scheme with the one provided. If +// an empty scheme is provided, OverrideScheme will return just the host:port. +func (a Address) OverrideScheme(scheme string) string { + addr := parse(a) + if scheme == "" { + return addr.hostPort() + } + return scheme + "://" + addr.hostPort() +} + +// WithScheme ensures that the string returned contains the scheme portion of a +// URL. Because an Address may not have a scheme (for example, it could be just +// "host:80"), this method can be applied to an address when it needs to be used +// as a URL. If the address's existing scheme is blank, the default scheme +// provided will be used. If address is blank, the default scheme will not be +// added; i.e. address will remain blank. +func (a Address) WithScheme(dflt string) string { + // If address is empty, don't add a scheme to it. + if a == "" { + return "" + } + + addr := parse(a) + if addr.scheme != "" { + return a.String() + } + return dflt + "://" + addr.hostPort() +} + +type addr struct { + scheme string + host string + port uint16 +} + +// parse breaks the address up into scheme://host:port. It currently assumes +// that very rigid structure; in other words, if an address does not follow that +// format, return values may be unexpected. +func parse(a Address) addr { + var scheme string + var host string + var port uint16 + + aStr := string(a) + + var hostPort string + if parts := strings.Split(aStr, "://"); len(parts) > 1 { + scheme = parts[0] + hostPort = parts[1] + } else { + hostPort = aStr + } + + if parts := strings.Split(hostPort, ":"); len(parts) == 2 { + host = parts[0] + portStr := parts[1] + port64, err := strconv.ParseInt(portStr, 10, 32) + if err == nil { + port = uint16(port64) + } + } else { + host = hostPort + } + + return addr{ + scheme: scheme, + host: host, + port: port, + } +} + +func (a addr) hostPort() string { + if a.port == 0 { + return a.host + } + return fmt.Sprintf("%s:%d", a.host, a.port) +} + +// AddressManager is an interface for any service which needs to maintain a list +// of addresses, and receive add/remove address requests from other services. +type AddressManager interface { + AddAddresses(context.Context, ...Address) error + RemoveAddresses(context.Context, ...Address) error +} + +// Ensure type implements interface. +var _ AddressManager = &NopAddressManager{} + +// NopAddressManager is a no-op implementation of the AddressManager interface. +type NopAddressManager struct{} + +func NewNopAddressManager() *NopAddressManager { + return &NopAddressManager{} +} + +func (a *NopAddressManager) AddAddresses(ctx context.Context, addrs ...Address) error { return nil } + +func (a *NopAddressManager) RemoveAddresses(ctx context.Context, addrs ...Address) error { return nil } diff --git a/dax/address_test.go b/dax/address_test.go new file mode 100644 index 000000000..f48985a91 --- /dev/null +++ b/dax/address_test.go @@ -0,0 +1,176 @@ +package dax_test + +import ( + "fmt" + "testing" + + "github.com/molecula/featurebase/v3/dax" + "github.com/stretchr/testify/assert" +) + +func TestAddress(t *testing.T) { + t.Run("Address", func(t *testing.T) { + tests := []struct { + addr dax.Address + expScheme string + expHostPort string + expHost string + expPort uint16 + }{ + { + // blank address + addr: "", + expScheme: "", + expHostPort: "", + expHost: "", + expPort: 0, + }, + { + // schema:// + addr: "http://", + expScheme: "http", + expHostPort: "", + expHost: "", + expPort: 0, + }, + { + // host + addr: "foo", + expScheme: "", + expHostPort: "foo", + expHost: "foo", + expPort: 0, + }, + { + // :port + addr: ":8080", + expScheme: "", + expHostPort: ":8080", + expHost: "", + expPort: 8080, + }, + { + // host:port + addr: "foo:8080", + expScheme: "", + expHostPort: "foo:8080", + expHost: "foo", + expPort: 8080, + }, + { + // schema://host:port + addr: "http://foo:8080", + expScheme: "http", + expHostPort: "foo:8080", + expHost: "foo", + expPort: 8080, + }, + { + // schema://host + addr: "http://foo", + expScheme: "http", + expHostPort: "foo", + expHost: "foo", + expPort: 0, + }, + { + // schema://:port + addr: "http://:8080", + expScheme: "http", + expHostPort: ":8080", + expHost: "", + expPort: 8080, + }, + { + // invalid port + addr: "http://foo:bar", + expScheme: "http", + expHostPort: "foo", + expHost: "foo", + expPort: 0, + }, + { + // :port outside of int16 range + addr: ":53308", + expScheme: "", + expHostPort: ":53308", + expHost: "", + expPort: 53308, + }, + } + + for i, test := range tests { + t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) { + assert.Equal(t, test.expScheme, test.addr.Scheme()) + assert.Equal(t, test.expHostPort, test.addr.HostPort()) + assert.Equal(t, test.expHost, test.addr.Host()) + assert.Equal(t, test.expPort, test.addr.Port()) + }) + } + }) + + t.Run("OverrideScheme", func(t *testing.T) { + tests := []struct { + addr dax.Address + scheme string + expAddr string + }{ + { + addr: "foo", + scheme: "http", + expAddr: "http://foo", + }, + { + addr: "http://foo", + scheme: "grpc", + expAddr: "grpc://foo", + }, + { + addr: "http://foo:8080", + scheme: "", + expAddr: "foo:8080", + }, + } + + for i, test := range tests { + t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) { + assert.Equal(t, test.expAddr, test.addr.OverrideScheme(test.scheme)) + }) + } + }) + + t.Run("WithScheme", func(t *testing.T) { + tests := []struct { + addr dax.Address + scheme string + expAddr string + }{ + { + addr: "foo", + scheme: "", + expAddr: "://foo", + }, + { + addr: "http://foo", + scheme: "grpc", + expAddr: "http://foo", + }, + { + addr: "http://foo:8080", + scheme: "", + expAddr: "http://foo:8080", + }, + { + addr: "foo:8080", + scheme: "grpc", + expAddr: "grpc://foo:8080", + }, + } + + for i, test := range tests { + t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) { + assert.Equal(t, test.expAddr, test.addr.WithScheme(test.scheme)) + }) + } + }) +} diff --git a/dax/boltdb/boltdb.go b/dax/boltdb/boltdb.go new file mode 100644 index 000000000..aef7a13c7 --- /dev/null +++ b/dax/boltdb/boltdb.go @@ -0,0 +1,151 @@ +// Package boltdb contains the boltdb implementations of the DAX interfaces. +package boltdb + +import ( + "context" + "os" + "path/filepath" + "strings" + "time" + + "github.com/molecula/featurebase/v3/errors" + bolt "go.etcd.io/bbolt" +) + +const ( + ErrFmtBucketNotFound = "boltdb: bucket '%s' not found" +) + +type Bucket []byte + +// DB represents the database connection. +type DB struct { + db *bolt.DB + ctx context.Context // background context + cancel func() // cancel background context + + // Datasource name. + DSN string + + // Destination for events to be published. + // EventService wtf.EventService + + // Returns the current time. Defaults to time.Now(). + // Can be mocked for tests. + Now func() time.Time + + filePath string + + // bucketQueue contains a list of buckets to create upon Open. + bucketQueue []Bucket +} + +// NewDB returns a new instance of DB associated with the given datasource name. +func NewDB(dsn string) *DB { + db := &DB{ + DSN: dsn, + Now: time.Now, + + //EventService: wtf.NopEventService(), + } + db.ctx, db.cancel = context.WithCancel(context.Background()) + return db +} + +// NewSvcBolt gets, opens, and creates buckets for a boltDB for a +// particular named service (the data file will be named after the +// service). +func NewSvcBolt(dir, svc string, buckets ...Bucket) (*DB, error) { + dir = strings.TrimPrefix(dir, "file:") + filename := filepath.Join(dir, svc+".boltdb") + db := NewDB("file:" + filename) + db.RegisterBuckets(buckets...) + err := db.Open() + return db, errors.Wrap(err, "opening") +} + +// path returns the file path to the boltdb database file. +func (db *DB) path() (string, error) { + if !strings.HasPrefix(db.DSN, "file:") { + return "", errors.New(errors.ErrUncoded, "boltdb package only supports a DSN beginning with `file:`") + } + + return db.DSN[5:], nil +} + +// RegisterBuckets queues up the buckets to be created when the database is +// first opened. +func (db *DB) RegisterBuckets(buckets ...Bucket) { + db.bucketQueue = append(db.bucketQueue, buckets...) +} + +// InitializeBuckets creates the given buckets if they do not already exist. +func (db *DB) InitializeBuckets(buckets ...Bucket) (err error) { + return db.db.Update(func(tx *bolt.Tx) error { + for _, bucket := range buckets { + if _, err := tx.CreateBucketIfNotExists(bucket); err != nil { + return errors.Wrapf(err, "creating bucket: %s", bucket) + } + } + return nil + }) +} + +// Open opens the database connection. +func (db *DB) Open() (err error) { + path, err := db.path() + if err != nil { + return errors.Wrap(err, "getting path from DSN") + } + + if err := os.MkdirAll(filepath.Dir(path), 0777); err != nil { + return errors.Wrapf(err, "mkdir %s", filepath.Dir(path)) + } else if db.db, err = bolt.Open(path, 0666, &bolt.Options{Timeout: 1 * time.Second}); err != nil { + return errors.Wrapf(err, "open file: %s", err) + } + + // cache the path in db.filePath. + db.filePath = path + + if err := db.InitializeBuckets(db.bucketQueue...); err != nil { + return errors.Wrap(err, "initializing buckets") + } + + // Reset the bucketQueue. + db.bucketQueue = make([]Bucket, 0) + + return nil +} + +// Close closes the database connection. +func (db *DB) Close() (err error) { + return db.db.Close() +} + +// BeginTx starts a transaction and returns a wrapper Tx type. This type +// provides a reference to the database and a fixed timestamp at the start of +// the transaction. The timestamp allows us to mock time during tests as well. +func (db *DB) BeginTx(ctx context.Context, writable bool) (*Tx, error) { + tx, err := db.db.Begin(writable) + if err != nil { + return nil, err + } + + // Return wrapper Tx that includes the transaction start time. + return &Tx{ + Tx: tx, + db: db, + now: db.Now().UTC().Truncate(time.Second), + }, nil +} + +// Tx wraps the SQL Tx object to provide a timestamp at the start of the transaction. +type Tx struct { + *bolt.Tx + db *DB + now time.Time +} + +func (db *DB) Path() string { + return db.filePath +} diff --git a/dax/boltdb/boltdb_test.go b/dax/boltdb/boltdb_test.go new file mode 100644 index 000000000..f284be42b --- /dev/null +++ b/dax/boltdb/boltdb_test.go @@ -0,0 +1,17 @@ +package boltdb_test + +import ( + "testing" + + "github.com/molecula/featurebase/v3/dax/test/boltdb" +) + +// Ensure the test database can open & close. +func TestDB(t *testing.T) { + db := boltdb.MustOpenDB(t) + defer boltdb.MustCloseDB(t, db) + + t.Cleanup(func() { + boltdb.CleanupDB(t, db.Path()) + }) +} diff --git a/dax/boltdb/directiveversion.go b/dax/boltdb/directiveversion.go new file mode 100644 index 000000000..56d140dd1 --- /dev/null +++ b/dax/boltdb/directiveversion.go @@ -0,0 +1,66 @@ +package boltdb + +import ( + "context" + "encoding/binary" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/errors" +) + +var ( + bucketDirective = Bucket("nodeDirective") + keyDirectiveVersion = []byte("directiveVersion") +) + +// DirectiveBuckets defines the buckets used by this package. It can be called +// during setup to create the buckets ahead of time. +var DirectiveBuckets []Bucket = []Bucket{ + bucketDirective, +} + +// Ensure type implements interface. +var _ dax.DirectiveVersion = (*DirectiveVersion)(nil) + +type DirectiveVersion struct { + db *DB +} + +func NewDirectiveVersion(db *DB) *DirectiveVersion { + return &DirectiveVersion{ + db: db, + } +} + +func (d *DirectiveVersion) Increment(ctx context.Context, delta uint64) (uint64, error) { + tx, err := d.db.BeginTx(ctx, true) + if err != nil { + return 0, errors.Wrap(err, "getting transaction") + } + defer tx.Rollback() + + bkt := tx.Bucket(bucketDirective) + if bkt == nil { + return 0, errors.Errorf(ErrFmtBucketNotFound, bucketDirective) + } + + var nextVersion uint64 = 1 // Start at 1; 0 is an invalid version. + + b := bkt.Get(keyDirectiveVersion) + if b != nil { + nextVersion = binary.LittleEndian.Uint64(b) + delta + } + + vsn := make([]byte, 8) + binary.LittleEndian.PutUint64(vsn, nextVersion) + + if err := bkt.Put(keyDirectiveVersion, vsn); err != nil { + return 0, errors.Wrap(err, "putting next directive version") + } + + if err := tx.Commit(); err != nil { + return 0, err + } + + return nextVersion, nil +} diff --git a/dax/boltdb/node.go b/dax/boltdb/node.go new file mode 100644 index 000000000..c0031213a --- /dev/null +++ b/dax/boltdb/node.go @@ -0,0 +1,157 @@ +package boltdb + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/errors" + "github.com/molecula/featurebase/v3/logger" +) + +var ( + bucketNodes = Bucket("nodeServiceNodes") +) + +// NodeServiceBuckets defines the buckets used by this package. It can be called +// during setup to create the buckets ahead of time. +var NodeServiceBuckets []Bucket = []Bucket{ + bucketNodes, +} + +// Ensure type implements interface. +var _ dax.NodeService = (*NodeService)(nil) + +// NodeService represents a service for managing nodes. +type NodeService struct { + db *DB + + logger logger.Logger +} + +// NewNodeService returns a new instance of NodeService with default values. +func NewNodeService(db *DB, logger logger.Logger) *NodeService { + return &NodeService{ + db: db, + logger: logger, + } +} + +func (s *NodeService) CreateNode(ctx context.Context, addr dax.Address, node *dax.Node) error { + tx, err := s.db.BeginTx(ctx, true) + if err != nil { + return errors.Wrap(err, "getting transaction") + } + defer tx.Rollback() + + bkt := tx.Bucket(bucketNodes) + if bkt == nil { + return errors.Errorf(ErrFmtBucketNotFound, bucketNodes) + } + + val, err := json.Marshal(node) + if err != nil { + return errors.Wrap(err, "marshalling node to json") + } + + if err := bkt.Put(addressKey(addr), val); err != nil { + return errors.Wrap(err, "putting node") + } + + return tx.Commit() +} + +func (s *NodeService) ReadNode(ctx context.Context, addr dax.Address) (*dax.Node, error) { + tx, err := s.db.BeginTx(ctx, false) + if err != nil { + return nil, errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + bkt := tx.Bucket(bucketNodes) + if bkt == nil { + return nil, errors.Errorf(ErrFmtBucketNotFound, bucketNodes) + } + + b := bkt.Get(addressKey(addr)) + if b == nil { + return nil, dax.NewErrNodeDoesNotExist(addr) + } + + node := &dax.Node{} + if err := json.Unmarshal(b, node); err != nil { + return nil, errors.Wrap(err, "unmarshalling node json") + } + + return node, nil +} + +func (s *NodeService) DeleteNode(ctx context.Context, addr dax.Address) error { + tx, err := s.db.BeginTx(ctx, true) + if err != nil { + return errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + bkt := tx.Bucket(bucketNodes) + if bkt == nil { + return errors.Errorf(ErrFmtBucketNotFound, bucketNodes) + } + + if err := bkt.Delete(addressKey(addr)); err != nil { + return errors.Wrapf(err, "deleting node key: %s", addressKey(addr)) + } + + return tx.Commit() +} + +func (s *NodeService) Nodes(ctx context.Context) ([]*dax.Node, error) { + tx, err := s.db.BeginTx(ctx, false) + if err != nil { + return nil, errors.Wrap(err, "getting tx") + } + defer tx.Rollback() + + nodes, err := s.getNodes(ctx, tx) + if err != nil { + return nil, errors.Wrap(err, "getting nodes") + } + + return nodes, nil +} + +func (s *NodeService) getNodes(ctx context.Context, tx *Tx) ([]*dax.Node, error) { + c := tx.Bucket(bucketNodes).Cursor() + + // Deserialize rows into Node objects. + nodes := make([]*dax.Node, 0) + + prefix := []byte(prefixFmtNodes) + for k, v := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, v = c.Next() { + if v == nil { + s.logger.Printf("nil value for key: %s", k) + continue + } + + node := &dax.Node{} + if err := json.Unmarshal(v, node); err != nil { + return nil, errors.Wrap(err, "unmarshalling node json") + } + + nodes = append(nodes, node) + } + + return nodes, nil +} + +const ( + prefixFmtNodes = "nodes/" +) + +// addressKey returns a key based on address. +func addressKey(addr dax.Address) []byte { + key := fmt.Sprintf(prefixFmtNodes+"%s", addr) + return []byte(key) +} diff --git a/dax/boltdb/node_test.go b/dax/boltdb/node_test.go new file mode 100644 index 000000000..6342d16ae --- /dev/null +++ b/dax/boltdb/node_test.go @@ -0,0 +1,55 @@ +package boltdb_test + +import ( + "context" + "testing" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/boltdb" + testbolt "github.com/molecula/featurebase/v3/dax/test/boltdb" + "github.com/molecula/featurebase/v3/errors" + "github.com/molecula/featurebase/v3/logger" + "github.com/stretchr/testify/assert" +) + +func TestNodeService(t *testing.T) { + db := testbolt.MustOpenDB(t) + defer testbolt.MustCloseDB(t, db) + + t.Cleanup(func() { + testbolt.CleanupDB(t, db.Path()) + }) + + ctx := context.Background() + + // Initialize the buckets. + assert.NoError(t, db.InitializeBuckets(boltdb.NodeServiceBuckets...)) + + t.Run("Nodes", func(t *testing.T) { + ns := boltdb.NewNodeService(db, logger.NopLogger) + + node1 := &dax.Node{ + Address: "localhost:10101", + RoleTypes: []dax.RoleType{ + "compute", + }, + } + + // Create node. + assert.NoError(t, ns.CreateNode(ctx, node1.Address, node1)) + + // Read node. + n, err := ns.ReadNode(ctx, node1.Address) + assert.NoError(t, err) + assert.Equal(t, node1, n) + + // Delete node. + assert.NoError(t, ns.DeleteNode(ctx, node1.Address)) + + // Read node. + _, err = ns.ReadNode(ctx, node1.Address) + if assert.Error(t, err) { + assert.True(t, errors.Is(err, dax.ErrNodeDoesNotExist)) + } + }) +} diff --git a/dax/boltdb/versionstore.go b/dax/boltdb/versionstore.go new file mode 100644 index 000000000..20171abd6 --- /dev/null +++ b/dax/boltdb/versionstore.go @@ -0,0 +1,768 @@ +package boltdb + +import ( + "bytes" + "context" + "encoding/binary" + "fmt" + "strconv" + "strings" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/inmem" + "github.com/molecula/featurebase/v3/errors" + "github.com/molecula/featurebase/v3/logger" +) + +var ( + bucketTables = Bucket("versionStoreTables") + bucketShards = Bucket("versionStoreShards") + bucketTableKeys = Bucket("versionStoreTableKeys") + bucketFieldKeys = Bucket("versionStoreFieldKeys") +) + +// VersionStoreBuckets defines the buckets used by this package. It can be +// called during setup to create the buckets ahead of time. +var VersionStoreBuckets []Bucket = []Bucket{ + bucketTables, + bucketShards, + bucketTableKeys, + bucketFieldKeys, +} + +// Ensure type implements interface. +var _ dax.VersionStore = (*VersionStore)(nil) + +// VersionStore manages all version info for shard, table keys, and field keys. +type VersionStore struct { + db *DB + + logger logger.Logger +} + +// NewVersionStore returns a new instance of VersionStore with default values. +func NewVersionStore(db *DB, logger logger.Logger) *VersionStore { + return &VersionStore{ + db: db, + logger: logger, + } +} + +func (s *VersionStore) AddTable(ctx context.Context, qtid dax.QualifiedTableID) error { + tx, err := s.db.BeginTx(ctx, true) + if err != nil { + return errors.Wrap(err, "getting transaction") + } + defer tx.Rollback() + + bkt := tx.Bucket(bucketTables) + if bkt == nil { + return errors.Errorf(ErrFmtBucketNotFound, bucketTables) + } + + if val := bkt.Get(tableKey(qtid)); val != nil { + return dax.NewErrTableIDExists(qtid) + } + + // The assumption is that we may store information about the table (other + // than just the fact that it exists). So for now, the value is an empty + // JSON object. + val := []byte("{}") + + if err := bkt.Put(tableKey(qtid), val); err != nil { + return errors.Wrap(err, "putting table") + } + + // Add the table to the "table index" of the other buckets. + // + // Shards + if bkt := tx.Bucket(bucketShards); bkt == nil { + return errors.Errorf(ErrFmtBucketNotFound, bucketShards) + } else if err := bkt.Put(tableKey(qtid), val); err != nil { + return errors.Wrap(err, "putting table into shards") + } + + // TableKeys. + if bkt := tx.Bucket(bucketTableKeys); bkt == nil { + return errors.Errorf(ErrFmtBucketNotFound, bucketTableKeys) + } else if err := bkt.Put(tableKey(qtid), val); err != nil { + return errors.Wrap(err, "putting table into table keys") + } + + // FieldKeys. + if bkt := tx.Bucket(bucketFieldKeys); bkt == nil { + return errors.Errorf(ErrFmtBucketNotFound, bucketFieldKeys) + } else if err := bkt.Put(tableKey(qtid), val); err != nil { + return errors.Wrap(err, "putting table into field keys") + } + + return tx.Commit() +} + +func (s *VersionStore) RemoveTable(ctx context.Context, qtid dax.QualifiedTableID) (dax.Shards, dax.Partitions, error) { + tx, err := s.db.BeginTx(ctx, true) + if err != nil { + return nil, nil, err + } + defer tx.Rollback() + + // Get the shards and partitions before deleting by table. + shards, err := s.getShards(ctx, tx, qtid) + if err != nil { + return nil, nil, err + } + + partitions, err := s.getPartitions(ctx, tx, qtid) + if err != nil { + return nil, nil, err + } + + if err := removeTable(ctx, tx, qtid); err != nil { + return nil, nil, err + } + + if err := tx.Commit(); err != nil { + return nil, nil, err + } + + return shards, partitions, nil +} + +func removeTable(ctx context.Context, tx *Tx, qtid dax.QualifiedTableID) error { + // Tables. + if bkt := tx.Bucket(bucketTables); bkt == nil { + return errors.Errorf(ErrFmtBucketNotFound, bucketTables) + } else if err := bkt.Delete(tableKey(qtid)); err != nil { + return errors.Wrap(err, "deleting table") + } + + // Shards. + if bkt := tx.Bucket(bucketShards); bkt == nil { + return errors.Errorf(ErrFmtBucketNotFound, bucketShards) + } else if err := bkt.Delete(tableKey(qtid)); err != nil { + return errors.Wrap(err, "deleting table in shards") + } else if err := deleteByPrefix(tx, bucketShards, []byte(fmt.Sprintf(prefixFmtShards, qtid.OrganizationID, qtid.DatabaseID, qtid.ID))); err != nil { + return errors.Wrap(err, "deleting shards for table") + } + + // TableKeys. + if bkt := tx.Bucket(bucketTableKeys); bkt == nil { + return errors.Errorf(ErrFmtBucketNotFound, bucketTableKeys) + } else if err := bkt.Delete(tableKey(qtid)); err != nil { + return errors.Wrap(err, "deleting table in table keys") + } else if err := deleteByPrefix(tx, bucketTableKeys, []byte(fmt.Sprintf(prefixFmtTableKeys, qtid.OrganizationID, qtid.DatabaseID, qtid.ID))); err != nil { + return errors.Wrap(err, "deleting table keys for table") + } + + // FieldKeys. + if bkt := tx.Bucket(bucketFieldKeys); bkt == nil { + return errors.Errorf(ErrFmtBucketNotFound, bucketFieldKeys) + } else if err := bkt.Delete(tableKey(qtid)); err != nil { + return errors.Wrap(err, "deleting table in field keys") + } else if err := deleteByPrefix(tx, bucketFieldKeys, []byte(fmt.Sprintf(prefixFmtFieldKeys, qtid.OrganizationID, qtid.DatabaseID, qtid.ID))); err != nil { + return errors.Wrap(err, "deleting field keys for table") + } + + return nil +} + +func deleteByPrefix(tx *Tx, bucket Bucket, prefix []byte) error { + bkt := tx.Bucket(bucket) + cursor := bkt.Cursor() + + // Deleting keys within the for loop seems to cause Next() to skip the next + // matching key because the Delete() call pops the item and effectively + // moves the cursor forward. Then calling Next() skips the item that was + // being pointed to after the delete. So, we're going to make a list of keys + // to delete, and then delete them outside of the cursor logic. + var keysToDelete [][]byte + for k, _ := cursor.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, _ = cursor.Next() { + keysToDelete = append(keysToDelete, k) + } + + for _, k := range keysToDelete { + if err := bkt.Delete(k); err != nil { + return errors.Wrapf(err, "deleting key: %s", k) + } + } + + return nil +} + +func (s *VersionStore) AddShards(ctx context.Context, qtid dax.QualifiedTableID, shards ...dax.Shard) error { + tx, err := s.db.BeginTx(ctx, true) + if err != nil { + return errors.Wrap(err, "getting transaction") + } + defer tx.Rollback() + + for _, shard := range shards { + if err := createShard(ctx, tx, qtid, shard); err != nil { + return errors.Wrap(err, "creating shard") + } + } + + return tx.Commit() +} + +func createShard(ctx context.Context, tx *Tx, qtid dax.QualifiedTableID, shard dax.Shard) error { + // TODO: validate data more formally + if shard.Version < 0 { + return errors.New(errors.ErrUncoded, fmt.Sprintf("invalid shard version: %d", shard.Version)) + } + + bkt := tx.Bucket(bucketShards) + if bkt == nil { + return errors.Errorf(ErrFmtBucketNotFound, bucketShards) + } + + // Ensure the table exists. + if val := bkt.Get(tableKey(qtid)); val == nil { + return dax.NewErrTableIDDoesNotExist(qtid) + } + + vsn := make([]byte, 8) + binary.LittleEndian.PutUint64(vsn, uint64(shard.Version)) + + return bkt.Put(shardKey(qtid, shard.Num), vsn) +} + +func (s *VersionStore) Shards(ctx context.Context, qtid dax.QualifiedTableID) (dax.Shards, bool, error) { + tx, err := s.db.BeginTx(ctx, false) + if err != nil { + return nil, false, errors.Wrap(err, "getting tx") + } + defer tx.Rollback() + + shards, err := s.getShards(ctx, tx, qtid) + if err != nil { + return nil, false, errors.Wrap(err, "getting shards") + } + + return shards, true, nil +} + +func (s *VersionStore) getShards(ctx context.Context, tx *Tx, qtid dax.QualifiedTableID) (dax.Shards, error) { + c := tx.Bucket(bucketShards).Cursor() + + // Deserialize rows into Shard objects. + shards := make(dax.Shards, 0) + + prefix := []byte(fmt.Sprintf(prefixFmtShards, qtid.OrganizationID, qtid.DatabaseID, qtid.ID)) + for k, v := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, v = c.Next() { + if v == nil { + s.logger.Printf("nil value for key: %s", k) + continue + } + + var shard dax.Shard + + shardNum, err := keyShardNum(k) + if err != nil { + return nil, errors.Wrapf(err, "getting shardNum from key: %v", k) + } + + shard.Num = shardNum + shard.Version = int(binary.LittleEndian.Uint64(v)) + + shards = append(shards, shard) + } + + return shards, nil +} + +// ShardVersion return the current version for the given table/shardNum. +// If a version is not being tracked, it returns a bool value of false. +func (s *VersionStore) ShardVersion(ctx context.Context, qtid dax.QualifiedTableID, shardNum dax.ShardNum) (int, bool, error) { + tx, err := s.db.BeginTx(ctx, false) + if err != nil { + return -1, false, err + } + defer tx.Rollback() + + return getShardVersion(ctx, tx, qtid, shardNum) +} + +func getShardVersion(ctx context.Context, tx *Tx, qtid dax.QualifiedTableID, shardNum dax.ShardNum) (int, bool, error) { + version := -1 + + bkt := tx.Bucket(bucketShards) + if bkt == nil { + return version, false, errors.Errorf(ErrFmtBucketNotFound, bucketShards) + } + + b := bkt.Get(shardKey(qtid, shardNum)) + if b == nil { + return version, false, nil + } + version = int(binary.LittleEndian.Uint64(b)) + + return version, true, nil +} + +func (s *VersionStore) ShardTables(ctx context.Context, qual dax.TableQualifier) (dax.TableIDs, error) { + tx, err := s.db.BeginTx(ctx, false) + if err != nil { + return nil, errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + return s.getTableIDs(ctx, tx, qual, bucketShards) +} + +func (s *VersionStore) getTableIDs(ctx context.Context, tx *Tx, qual dax.TableQualifier, bucket Bucket) (dax.TableIDs, error) { + c := tx.Bucket(bucket).Cursor() + + // Deserialize rows into Tables objects. + tableIDs := make(dax.TableIDs, 0) + + prefix := []byte(fmt.Sprintf(prefixFmtTables, qual.OrganizationID, qual.DatabaseID)) + for k, v := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, v = c.Next() { + if v == nil { + s.logger.Printf("nil value for key: %s", k) + continue + } + + var tableID dax.TableID + + tableID, err := keyTableID(k) + if err != nil { + return nil, errors.Wrapf(err, "getting table name from key: %v", k) + } + + tableIDs = append(tableIDs, tableID) + } + + return tableIDs, nil +} + +func (s *VersionStore) bucketTables(ctx context.Context, bucket Bucket) ([]dax.QualifiedTableID, error) { + tx, err := s.db.BeginTx(ctx, false) + if err != nil { + return nil, errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + c := tx.Bucket(bucket).Cursor() + + // Deserialize rows into Tables objects. + qtids := make([]dax.QualifiedTableID, 0) + + prefix := []byte(prefixTables) + for k, v := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, v = c.Next() { + if v == nil { + s.logger.Printf("nil value for key: %s", k) + continue + } + + qtid, err := keyQualifiedTableID(k) + if err != nil { + return nil, errors.Wrapf(err, "getting qualified table id from key: %v", k) + } + + qtids = append(qtids, qtid) + } + + return qtids, nil +} + +// AddPartitions adds new partitions to be managed by VersionStore. It returns +// the number of partitions added or an error. +func (s *VersionStore) AddPartitions(ctx context.Context, qtid dax.QualifiedTableID, partitions ...dax.Partition) error { + tx, err := s.db.BeginTx(ctx, true) + if err != nil { + return errors.Wrap(err, "getting transaction") + } + defer tx.Rollback() + + for _, partition := range partitions { + if err := createPartition(ctx, tx, qtid, partition); err != nil { + return errors.Wrap(err, "creating partition") + } + } + + return tx.Commit() +} + +func createPartition(ctx context.Context, tx *Tx, qtid dax.QualifiedTableID, partition dax.Partition) error { + // TODO: validate data more formally + if partition.Version < 0 { + return errors.New(errors.ErrUncoded, fmt.Sprintf("invalid partition version: %d", partition.Version)) + } + + bkt := tx.Bucket(bucketTableKeys) + if bkt == nil { + return errors.Errorf(ErrFmtBucketNotFound, bucketTableKeys) + } + + // Ensure the table exists. + if val := bkt.Get(tableKey(qtid)); val == nil { + return dax.NewErrTableIDDoesNotExist(qtid) + } + + vsn := make([]byte, 8) + binary.LittleEndian.PutUint64(vsn, uint64(partition.Version)) + + return bkt.Put(partitionKey(qtid, partition.Num), vsn) +} + +func (s *VersionStore) Partitions(ctx context.Context, qtid dax.QualifiedTableID) (dax.Partitions, bool, error) { + tx, err := s.db.BeginTx(ctx, false) + if err != nil { + return nil, false, errors.Wrap(err, "getting tx") + } + defer tx.Rollback() + + partitions, err := s.getPartitions(ctx, tx, qtid) + if err != nil { + return nil, false, errors.Wrap(err, "getting partitions") + } + + return partitions, true, nil +} + +func (s *VersionStore) getPartitions(ctx context.Context, tx *Tx, qtid dax.QualifiedTableID) (dax.Partitions, error) { + c := tx.Bucket(bucketTableKeys).Cursor() + + // Deserialize rows into Partition objects. + partitions := make(dax.Partitions, 0) + + prefix := []byte(fmt.Sprintf(prefixFmtTableKeys, qtid.OrganizationID, qtid.DatabaseID, qtid.ID)) + for k, v := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, v = c.Next() { + if v == nil { + s.logger.Printf("nil value for key: %s", k) + continue + } + + var partition dax.Partition + + partitionNum, err := keyPartitionNum(k) + if err != nil { + return nil, errors.Wrapf(err, "getting partitionNum from key: %v", k) + } + + partition.Num = partitionNum + partition.Version = int(binary.LittleEndian.Uint64(v)) + + partitions = append(partitions, partition) + } + + return partitions, nil +} + +func (s *VersionStore) PartitionVersion(ctx context.Context, qtid dax.QualifiedTableID, partitionNum dax.PartitionNum) (int, bool, error) { + tx, err := s.db.BeginTx(ctx, false) + if err != nil { + return -1, false, err + } + defer tx.Rollback() + + return getPartitionVersion(ctx, tx, qtid, partitionNum) +} + +func getPartitionVersion(ctx context.Context, tx *Tx, qtid dax.QualifiedTableID, partitionNum dax.PartitionNum) (int, bool, error) { + version := -1 + + bkt := tx.Bucket(bucketTableKeys) + if bkt == nil { + return version, false, errors.Errorf(ErrFmtBucketNotFound, bucketTableKeys) + } + + b := bkt.Get(partitionKey(qtid, partitionNum)) + if b == nil { + return version, false, nil + } + version = int(binary.LittleEndian.Uint64(b)) + + return version, true, nil +} + +func (s *VersionStore) PartitionTables(ctx context.Context, qual dax.TableQualifier) (dax.TableIDs, error) { + tx, err := s.db.BeginTx(ctx, false) + if err != nil { + return nil, errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + return s.getTableIDs(ctx, tx, qual, bucketTableKeys) +} + +// AddFields adds new fields to be managed by VersionStore. It returns the +// number of fields added or an error. +func (s *VersionStore) AddFields(ctx context.Context, qtid dax.QualifiedTableID, fields ...dax.FieldVersion) error { + tx, err := s.db.BeginTx(ctx, true) + if err != nil { + return err + } + defer tx.Rollback() + + for _, field := range fields { + if err := createFieldVersion(ctx, tx, qtid, field); err != nil { + return errors.Wrap(err, "creating field version") + } + } + + return tx.Commit() +} + +func createFieldVersion(ctx context.Context, tx *Tx, qtid dax.QualifiedTableID, field dax.FieldVersion) error { + // TODO: validate data more formally + if field.Version < 0 { + return errors.New(errors.ErrUncoded, fmt.Sprintf("invalid field version: %d", field.Version)) + } + + bkt := tx.Bucket(bucketFieldKeys) + if bkt == nil { + return errors.Errorf(ErrFmtBucketNotFound, bucketFieldKeys) + } + + // Ensure the table exists. + if val := bkt.Get(tableKey(qtid)); val == nil { + return dax.NewErrTableIDDoesNotExist(qtid) + } + + vsn := make([]byte, 8) + binary.LittleEndian.PutUint64(vsn, uint64(field.Version)) + + return bkt.Put(fieldKey(qtid, field.Name), vsn) +} + +func (s *VersionStore) Fields(ctx context.Context, qtid dax.QualifiedTableID) (dax.FieldVersions, bool, error) { + tx, err := s.db.BeginTx(ctx, false) + if err != nil { + return nil, false, errors.Wrap(err, "getting tx") + } + defer tx.Rollback() + + fields, err := s.getFields(ctx, tx, qtid) + if err != nil { + return nil, false, errors.Wrap(err, "getting fields") + } + + return fields, true, nil +} + +func (s *VersionStore) getFields(ctx context.Context, tx *Tx, qtid dax.QualifiedTableID) (dax.FieldVersions, error) { + c := tx.Bucket(bucketFieldKeys).Cursor() + + // Deserialize rows into FieldVersion objects. + fieldVersions := make(dax.FieldVersions, 0) + + prefix := []byte(fmt.Sprintf(prefixFmtFieldKeys, qtid.OrganizationID, qtid.DatabaseID, qtid.ID)) + for k, v := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, v = c.Next() { + if v == nil { + s.logger.Printf("nil value for key: %s", k) + continue + } + + var fieldVersion dax.FieldVersion + + fieldName, err := keyFieldName(k) + if err != nil { + return nil, errors.Wrapf(err, "getting partitionNum from key: %v", k) + } + + fieldVersion.Name = fieldName + fieldVersion.Version = int(binary.LittleEndian.Uint64(v)) + + fieldVersions = append(fieldVersions, fieldVersion) + } + + return fieldVersions, nil +} + +func (s *VersionStore) FieldVersion(ctx context.Context, qtid dax.QualifiedTableID, field dax.FieldName) (int, bool, error) { + tx, err := s.db.BeginTx(ctx, false) + if err != nil { + return -1, false, err + } + defer tx.Rollback() + + return getFieldVersion(ctx, tx, qtid, field) +} + +func getFieldVersion(ctx context.Context, tx *Tx, qtid dax.QualifiedTableID, field dax.FieldName) (int, bool, error) { + version := -1 + + bkt := tx.Bucket(bucketFieldKeys) + if bkt == nil { + return version, false, errors.Errorf(ErrFmtBucketNotFound, bucketFieldKeys) + } + + b := bkt.Get(fieldKey(qtid, field)) + if b == nil { + return version, false, nil + } + version = int(binary.LittleEndian.Uint64(b)) + + return version, true, nil +} + +func (s *VersionStore) FieldTables(ctx context.Context, qual dax.TableQualifier) (dax.TableIDs, error) { + tx, err := s.db.BeginTx(ctx, false) + if err != nil { + return nil, errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + return s.getTableIDs(ctx, tx, qual, bucketFieldKeys) +} + +// Copy returns an in-memory copy of VersionStore. +func (s *VersionStore) Copy(ctx context.Context) (dax.VersionStore, error) { + new := inmem.NewVersionStore() + + // shards. + qtids, err := s.bucketTables(ctx, bucketShards) + if err != nil { + return nil, errors.Wrap(err, "getting shard tables") + } + for _, qtid := range qtids { + shards, found, err := s.Shards(ctx, qtid) + if err != nil { + return nil, errors.Wrap(err, "getting shards") + } else if !found { + continue + } + _ = new.AddTable(ctx, qtid) + new.AddShards(ctx, qtid, shards...) + } + + // tableKeys. + qtids, err = s.bucketTables(ctx, bucketTableKeys) + if err != nil { + return nil, errors.Wrap(err, "getting table key tables") + } + for _, qtid := range qtids { + partitions, found, err := s.Partitions(ctx, qtid) + if err != nil { + return nil, errors.Wrap(err, "getting partitions") + } else if !found { + continue + } + _ = new.AddTable(ctx, qtid) + new.AddPartitions(ctx, qtid, partitions...) + } + + // fieldKeys. + qtids, err = s.bucketTables(ctx, bucketFieldKeys) + if err != nil { + return nil, errors.Wrap(err, "getting field key tables") + } + for _, qtid := range qtids { + fields, found, err := s.Fields(ctx, qtid) + if err != nil { + return nil, errors.Wrap(err, "getting fields") + } else if !found { + continue + } + _ = new.AddTable(ctx, qtid) + new.AddFields(ctx, qtid, fields...) + } + + return new, nil +} + +///////////////////////////////////////////////////////// + +const ( + prefixShards = "shards/" + prefixFmtShards = prefixShards + "%s/%s/%s/" + + prefixTableKeys = "tablekeys/" + prefixFmtTableKeys = prefixTableKeys + "%s/%s/%s/" + + prefixFieldKeys = "fieldkeys/" + prefixFmtFieldKeys = prefixFieldKeys + "%s/%s/%s/" + + prefixTables = "tables/" + prefixFmtTables = prefixTables + "%s/%s/" +) + +// tableKey returns a key based on table name. +func tableKey(qtid dax.QualifiedTableID) []byte { + qual := qtid.TableQualifier + key := fmt.Sprintf(prefixFmtTables+"%s", qual.OrganizationID, qual.DatabaseID, qtid.ID) + return []byte(key) +} + +// keyTableID gets the table ID out of the key. +func keyTableID(key []byte) (dax.TableID, error) { + parts := strings.Split(string(key), "/") + if len(parts) != 4 { + return "", errors.New(errors.ErrUncoded, "table key format expected: `tables/orgID/dbID/tableID`") + } + + return dax.TableID(parts[3]), nil +} + +// keyQualifiedTableID gets the qualified table ID out of the key. +func keyQualifiedTableID(key []byte) (dax.QualifiedTableID, error) { + parts := strings.Split(string(key), "/") + if len(parts) != 4 { + return dax.QualifiedTableID{}, errors.New(errors.ErrUncoded, "table key format expected: `tables/orgID/dbID/tableID`") + } + + return dax.NewQualifiedTableID( + dax.NewTableQualifier(dax.OrganizationID(parts[1]), dax.DatabaseID(parts[2])), + dax.TableID(parts[3]), + ), nil +} + +// shardKey returns a key based on table and shard. +func shardKey(qtid dax.QualifiedTableID, shard dax.ShardNum) []byte { + key := fmt.Sprintf(prefixFmtShards+"%d", qtid.OrganizationID, qtid.DatabaseID, qtid.ID, shard) + return []byte(key) +} + +// keyShardNum gets the shardNum out of the key. +func keyShardNum(key []byte) (dax.ShardNum, error) { + parts := strings.Split(string(key), "/") + if len(parts) != 5 { + return 0, errors.New(errors.ErrUncoded, "shard key format expected: `shards/orgID/dbID/table/shard`") + } + + intVar, err := strconv.Atoi(parts[4]) + if err != nil { + return 0, errors.Wrapf(err, "converting string to shardNum: %s", parts[4]) + } + + return dax.ShardNum(intVar), nil +} + +// partitionKey returns a key based on table and partition. +func partitionKey(qtid dax.QualifiedTableID, partition dax.PartitionNum) []byte { + key := fmt.Sprintf(prefixFmtTableKeys+"%d", qtid.OrganizationID, qtid.DatabaseID, qtid.ID, partition) + return []byte(key) +} + +// keyPartitionNum gets the partitionNum out of the key. +func keyPartitionNum(key []byte) (dax.PartitionNum, error) { + parts := strings.Split(string(key), "/") + if len(parts) != 5 { + return 0, errors.New(errors.ErrUncoded, "partition key format expected: `tablekeys/orgID/dbID/table/partition`") + } + + intVar, err := strconv.Atoi(parts[4]) + if err != nil { + return 0, errors.Wrapf(err, "converting string to partitionNum: %s", parts[4]) + } + + return dax.PartitionNum(intVar), nil +} + +// fieldKey returns a key based on table and field. +func fieldKey(qtid dax.QualifiedTableID, field dax.FieldName) []byte { + key := fmt.Sprintf(prefixFmtFieldKeys+"%s", qtid.OrganizationID, qtid.DatabaseID, qtid.ID, field) + return []byte(key) +} + +// keyFieldName gets the fieldName out of the key. +func keyFieldName(key []byte) (dax.FieldName, error) { + parts := strings.Split(string(key), "/") + if len(parts) != 5 { + return "", errors.New(errors.ErrUncoded, "field key format expected: `fieldkeys/orgID/dbID/table/field`") + } + + return dax.FieldName(parts[4]), nil +} diff --git a/dax/boltdb/versionstore_test.go b/dax/boltdb/versionstore_test.go new file mode 100644 index 000000000..cead07425 --- /dev/null +++ b/dax/boltdb/versionstore_test.go @@ -0,0 +1,388 @@ +package boltdb_test + +import ( + "context" + "fmt" + "sort" + "testing" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/boltdb" + testbolt "github.com/molecula/featurebase/v3/dax/test/boltdb" + "github.com/molecula/featurebase/v3/logger" + "github.com/stretchr/testify/assert" +) + +func TestVersionStore(t *testing.T) { + db := testbolt.MustOpenDB(t) + defer testbolt.MustCloseDB(t, db) + + ctx := context.Background() + + t.Cleanup(func() { + testbolt.CleanupDB(t, db.Path()) + }) + + orgID := dax.OrganizationID("acme") + dbID := dax.DatabaseID("db1") + + qual := dax.NewTableQualifier(orgID, dbID) + + // Initialize the buckets. + assert.NoError(t, db.InitializeBuckets(boltdb.VersionStoreBuckets...)) + + t.Run("Tables", func(t *testing.T) { + vs := boltdb.NewVersionStore(db, logger.NopLogger) + + qtids := newQualifiedTableIDs(t, qual, 3) + qtid1 := qtids[0] + qtid2 := qtids[1] + qtid3 := qtids[2] + defer vs.RemoveTable(ctx, qtid1) + defer vs.RemoveTable(ctx, qtid2) + defer vs.RemoveTable(ctx, qtid3) + + // Add table 1. + assert.NoError(t, vs.AddTable(ctx, qtid1)) + + // Add table 2. + assert.NoError(t, vs.AddTable(ctx, qtid2)) + + // Add table 3. + assert.NoError(t, vs.AddTable(ctx, qtid3)) + }) + + t.Run("Shards", func(t *testing.T) { + vs := boltdb.NewVersionStore(db, logger.NopLogger) + + qtids := newQualifiedTableIDs(t, qual, 3) + qtid1 := qtids[0] + qtid2 := qtids[1] + qtid3 := qtids[2] + + // Add tables. + assert.NoError(t, vs.AddTable(ctx, qtid1)) + assert.NoError(t, vs.AddTable(ctx, qtid2)) + assert.NoError(t, vs.AddTable(ctx, qtid3)) + defer vs.RemoveTable(ctx, qtid1) + defer vs.RemoveTable(ctx, qtid2) + defer vs.RemoveTable(ctx, qtid3) + + // Create some shards to insert into the table. + shards := make(dax.Shards, 3) + for i := range shards { + shards[i] = dax.Shard{ + Num: dax.ShardNum(i), + Version: i * 2, + } + } + + // Add shards to table 1. + { + err := vs.AddShards(ctx, qtid1, shards...) + assert.NoError(t, err) + } + + // Add shards to table 2. + { + err := vs.AddShards(ctx, qtid2, shards...) + assert.NoError(t, err) + } + + // Fetch a shard and compare. + { + version, found, err := vs.ShardVersion(ctx, qtid1, 2) + assert.NoError(t, err) + assert.True(t, found) + assert.Equal(t, 4, version) + } + + // Fetch all shards and compare. + { + shrds, found, err := vs.Shards(ctx, qtid1) + assert.NoError(t, err) + assert.True(t, found) + assert.Equal(t, shards, shrds) + } + + // Fetch tables. + { + tblIDs, err := vs.ShardTables(ctx, qual) + assert.NoError(t, err) + exp := dax.TableIDs{qtid1.ID, qtid2.ID, qtid3.ID} + assert.Equal(t, exp, tblIDs) + } + + // Remove table 1. + { + shards, partitions, err := vs.RemoveTable(ctx, qtid1) + assert.NoError(t, err) + assert.Equal(t, shards, shards) + assert.Equal(t, dax.Partitions{}, partitions) + } + + // Fetch all shards and compare. + { + shrds, found, err := vs.Shards(ctx, qtid1) + assert.NoError(t, err) + assert.True(t, found) + assert.Equal(t, dax.Shards{}, shrds) + } + }) + + t.Run("Partitions", func(t *testing.T) { + vs := boltdb.NewVersionStore(db, logger.NopLogger) + + // Create some partitions to insert into the table. + partitions := make(dax.Partitions, 3) + for i := range partitions { + partitions[i] = dax.Partition{ + Num: dax.PartitionNum(i), + Version: i * 2, + } + } + + qtids := newQualifiedTableIDs(t, qual, 2) + qtid1 := qtids[0] + qtid2 := qtids[1] + + // Add tables. + assert.NoError(t, vs.AddTable(ctx, qtid1)) + assert.NoError(t, vs.AddTable(ctx, qtid2)) + defer vs.RemoveTable(ctx, qtid1) + defer vs.RemoveTable(ctx, qtid2) + + // Add partitions to table 1. + { + err := vs.AddPartitions(ctx, qtid1, partitions...) + assert.NoError(t, err) + } + + // Add partitions to table 2. + { + err := vs.AddPartitions(ctx, qtid2, partitions...) + assert.NoError(t, err) + } + + // Fetch a partition and compare. + { + version, found, err := vs.PartitionVersion(ctx, qtid1, 2) + assert.NoError(t, err) + assert.True(t, found) + assert.Equal(t, 4, version) + } + + // Fetch all partitions and compare. + { + parts, found, err := vs.Partitions(ctx, qtid1) + assert.NoError(t, err) + assert.True(t, found) + assert.Equal(t, partitions, parts) + } + + // Fetch tables. + { + tblIDs, err := vs.PartitionTables(ctx, qual) + assert.NoError(t, err) + exp := dax.TableIDs{qtid1.ID, qtid2.ID} + assert.Equal(t, exp, tblIDs) + } + + // Remove table 1. + { + shards, partitions, err := vs.RemoveTable(ctx, qtid1) + assert.NoError(t, err) + assert.Equal(t, dax.Shards{}, shards) + assert.Equal(t, partitions, partitions) + } + + // Fetch all partitions and compare. + { + parts, found, err := vs.Partitions(ctx, qtid1) + assert.NoError(t, err) + assert.True(t, found) + assert.Equal(t, dax.Partitions{}, parts) + } + }) + + t.Run("FieldVersions", func(t *testing.T) { + vs := boltdb.NewVersionStore(db, logger.NopLogger) + + qtids := newQualifiedTableIDs(t, qual, 2) + qtid1 := qtids[0] + qtid2 := qtids[1] + + // Add tables. + assert.NoError(t, vs.AddTable(ctx, qtid1)) + assert.NoError(t, vs.AddTable(ctx, qtid2)) + defer vs.RemoveTable(ctx, qtid1) + defer vs.RemoveTable(ctx, qtid2) + + // Create some fieldVersions to insert into the table. + fieldVersions := make(dax.FieldVersions, 3) + for i := range fieldVersions { + fieldVersions[i] = dax.FieldVersion{ + Name: dax.FieldName(fmt.Sprintf("fld-%d", i)), + Version: i * 2, + } + } + + // Add fieldVersions to table 1. + { + err := vs.AddFields(ctx, qtid1, fieldVersions...) + assert.NoError(t, err) + } + + // Add fieldVersions to table 2. + { + err := vs.AddFields(ctx, qtid2, fieldVersions...) + assert.NoError(t, err) + } + + // Fetch a fieldVersion and compare. + { + version, found, err := vs.FieldVersion(ctx, qtid1, dax.FieldName("fld-2")) + assert.NoError(t, err) + assert.True(t, found) + assert.Equal(t, 4, version) + } + + // Fetch all fieldVersions and compare. + { + flds, found, err := vs.Fields(ctx, qtid1) + assert.NoError(t, err) + assert.True(t, found) + assert.Equal(t, fieldVersions, flds) + } + + // Fetch tables. + { + tblIDs, err := vs.FieldTables(ctx, qual) + assert.NoError(t, err) + exp := dax.TableIDs{qtid1.ID, qtid2.ID} + assert.Equal(t, exp, tblIDs) + } + + // Remove table 1. + { + shards, partitions, err := vs.RemoveTable(ctx, qtid1) + assert.NoError(t, err) + assert.Equal(t, dax.Shards{}, shards) + assert.Equal(t, dax.Partitions{}, partitions) + } + + // Fetch all fieldVersions and compare. + { + flds, found, err := vs.Fields(ctx, qtid1) + assert.NoError(t, err) + assert.True(t, found) + assert.Equal(t, dax.FieldVersions{}, flds) + } + }) + + t.Run("Copy", func(t *testing.T) { + vs := boltdb.NewVersionStore(db, logger.NopLogger) + + qtids := newQualifiedTableIDs(t, qual, 1) + qtid1 := qtids[0] + + // Add tables. + assert.NoError(t, vs.AddTable(ctx, qtid1)) + defer vs.RemoveTable(ctx, qtid1) + + // Create some shards to insert into the table. + shards := make(dax.Shards, 3) + for i := range shards { + shards[i] = dax.Shard{ + Num: dax.ShardNum(i), + Version: i * 2, + } + } + + // Create some partitions to insert into the table. + partitions := make(dax.Partitions, 3) + for i := range partitions { + partitions[i] = dax.Partition{ + Num: dax.PartitionNum(i), + Version: i * 2, + } + } + + // Create some fieldVersions to insert into the table. + fieldVersions := make(dax.FieldVersions, 3) + for i := range fieldVersions { + fieldVersions[i] = dax.FieldVersion{ + Name: dax.FieldName(fmt.Sprintf("fld-%d", i)), + Version: i * 2, + } + } + + // Add shards to table 1. + { + err := vs.AddShards(ctx, qtid1, shards...) + assert.NoError(t, err) + } + + // Add partitions to table 1. + { + err := vs.AddPartitions(ctx, qtid1, partitions...) + assert.NoError(t, err) + } + + // Add fieldVersions to table 1. + { + err := vs.AddFields(ctx, qtid1, fieldVersions...) + assert.NoError(t, err) + } + + copy, err := vs.Copy(ctx) + assert.NoError(t, err) + + // Fetch a shard and compare. + { + version, found, err := copy.ShardVersion(ctx, qtid1, 2) + assert.NoError(t, err) + assert.True(t, found) + assert.Equal(t, 4, version) + } + + // Fetch all partitions and compare. + { + parts, found, err := copy.Partitions(ctx, qtid1) + assert.NoError(t, err) + assert.True(t, found) + assert.Equal(t, partitions, parts) + } + + // Fetch all fieldVersions and compare. + { + flds, found, err := copy.Fields(ctx, qtid1) + assert.NoError(t, err) + assert.True(t, found) + assert.Equal(t, fieldVersions, flds) + } + }) +} + +// newQualifiedTableIDs is a test helper function which generates a slice of n +// qtid. The entries in the slice will be ordered by TableID. +func newQualifiedTableIDs(t *testing.T, qual dax.TableQualifier, n int) []dax.QualifiedTableID { + t.Helper() + + qtids := make([]dax.QualifiedTableID, n) + for i := range qtids { + tbl := dax.NewTable("testvstore") + tbl.CreateID() + qtids[i] = dax.NewQualifiedTableID( + qual, + tbl.ID, + ) + } + + // sort the qtids by ID + sort.Slice(qtids, func(i, j int) bool { + return qtids[i].ID < qtids[j].ID + }) + + return qtids +} diff --git a/dax/computer/alpha/alpha.go b/dax/computer/alpha/alpha.go new file mode 100644 index 000000000..c8f8bfd1a --- /dev/null +++ b/dax/computer/alpha/alpha.go @@ -0,0 +1,24 @@ +package alpha + +import ( + "fmt" + "path" + + "github.com/molecula/featurebase/v3/dax" +) + +const ( + keysFileName = "keys" +) + +func partitionBucket(table dax.TableKey, partition dax.PartitionNum) string { + return path.Join(string(table), "partition", fmt.Sprintf("%d", partition)) +} + +func shardKey(shard dax.ShardNum) string { + return path.Join("shard", fmt.Sprintf("%d", shard)) +} + +func fieldBucket(table dax.TableKey, field dax.FieldName) string { + return path.Join(string(table), "field", string(field)) +} diff --git a/dax/computer/alpha/snapshot.go b/dax/computer/alpha/snapshot.go new file mode 100644 index 000000000..5da02f213 --- /dev/null +++ b/dax/computer/alpha/snapshot.go @@ -0,0 +1,104 @@ +// Package alpha contains an implementation of the SnapshotReadWriter interface. +// In the case where a sub-service (such as snapshotter) implements these +// interfaces directly with both its service and its http client, then we don't +// need this middle implementation layer. But in this case, the Snapshotter +// operates as a third-party service might, meaning its API methods don't align +// with what FeatureBase needs to call. So this implementation acts as a +// translation later between the featurebase-to-snapshotter interface, and the +// third-party Snapshotter service. +package alpha + +import ( + "context" + "io" + + featurebase "github.com/molecula/featurebase/v3" + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/computer" + "github.com/molecula/featurebase/v3/errors" +) + +// Ensure type implements interface. +var _ computer.SnapshotReadWriter = &alphaSnapshot{} + +// alphaSnapshot uses a Snapshotter implementation (which could be, for +// example, an http client or a locally running sub-service) to store its +// snapshots. +type alphaSnapshot struct { + ss featurebase.Snapshotter +} + +func NewAlphaSnapshot(sser featurebase.Snapshotter) *alphaSnapshot { + return &alphaSnapshot{ + ss: sser, + } +} + +func (s *alphaSnapshot) WriteShardData(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, shard dax.ShardNum, version int, rc io.ReadCloser) error { + bucket := partitionBucket(qtid.Key(), partition) + key := shardKey(shard) + + if err := s.ss.Write(bucket, key, version, rc); err != nil { + return errors.Wrapf(err, "writing shard data: %s, %d", key, version) + } + + return nil +} + +func (s *alphaSnapshot) ReadShardData(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, shard dax.ShardNum, version int) (io.ReadCloser, error) { + bucket := partitionBucket(qtid.Key(), partition) + key := shardKey(shard) + + rc, err := s.ss.Read(bucket, key, version) + if err != nil { + return nil, errors.Wrapf(err, "reading shard data: %s, %s, %d", bucket, key, version) + } + + return rc, nil +} + +func (s *alphaSnapshot) WriteTableKeys(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, version int, wrTo io.WriterTo) error { + bucket := partitionBucket(qtid.Key(), partition) + key := keysFileName + + if err := s.ss.WriteTo(bucket, key, version, wrTo); err != nil { + return errors.Wrapf(err, "writing table keys: %s, %d", key, version) + } + + return nil +} + +func (s *alphaSnapshot) ReadTableKeys(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, version int) (io.ReadCloser, error) { + bucket := partitionBucket(qtid.Key(), partition) + key := keysFileName + + rc, err := s.ss.Read(bucket, key, version) + if err != nil { + return nil, errors.Wrapf(err, "reading table keys: %s, %s, %d", bucket, key, version) + } + + return rc, nil +} + +func (s *alphaSnapshot) WriteFieldKeys(ctx context.Context, qtid dax.QualifiedTableID, field dax.FieldName, version int, wrTo io.WriterTo) error { + bucket := fieldBucket(qtid.Key(), field) + key := keysFileName + + if err := s.ss.WriteTo(bucket, key, version, wrTo); err != nil { + return errors.Wrapf(err, "writing field keys: %s, %d", key, version) + } + + return nil +} + +func (s *alphaSnapshot) ReadFieldKeys(ctx context.Context, qtid dax.QualifiedTableID, field dax.FieldName, version int) (io.ReadCloser, error) { + bucket := fieldBucket(qtid.Key(), field) + key := keysFileName + + rc, err := s.ss.Read(bucket, key, version) + if err != nil { + return nil, errors.Wrapf(err, "reading field keys: %s, %s, %d", bucket, key, version) + } + + return rc, nil +} diff --git a/dax/computer/alpha/writelog.go b/dax/computer/alpha/writelog.go new file mode 100644 index 000000000..cb87a7b62 --- /dev/null +++ b/dax/computer/alpha/writelog.go @@ -0,0 +1,331 @@ +// Package alpha contains an implementation of the WriteLogReader and +// WriteLogWriter interfaces. In the case where a sub-service (such as +// writelogger) implements these interfaces directly with both its service and +// its http client, then we don't need this middle implementation layer. But in +// this case, the WriteLogger operates as a third-party service might, meaning +// its API methods don't align with what FeatureBase needs to call. So this +// implementation acts as a translation later between the +// featurebase-to-writelogger interface, and the third-party WriteLogger +// service. +package alpha + +import ( + "bufio" + "context" + "encoding/json" + "io" + + featurebase "github.com/molecula/featurebase/v3" + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/computer" + "github.com/molecula/featurebase/v3/errors" +) + +// Ensure type implements interface. +var _ computer.WriteLogReader = &alphaWriteLog{} +var _ computer.WriteLogWriter = &alphaWriteLog{} + +// alphaWriteLog uses a WLer implementation (which could be, for example, an +// http client or a locally running sub-service) to store its log messages. +type alphaWriteLog struct { + wl featurebase.WriteLogger +} + +func NewAlphaWriteLog(wler featurebase.WriteLogger) *alphaWriteLog { + return &alphaWriteLog{ + wl: wler, + } +} + +func (w *alphaWriteLog) CreateTableKeys(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, version int, m map[string]uint64) error { + msg := computer.PartitionKeyMap{ + TableKey: qtid.Key(), + Partition: partition, + StringToID: m, + } + + b, err := json.Marshal(msg) + if err != nil { + return errors.Wrap(err, "marshalling partition key map to json") + } + + bucket := partitionBucket(qtid.Key(), partition) + + if err := w.wl.AppendMessage(bucket, keysFileName, version, b); err != nil { + return errors.Wrapf(err, "appending partition key message: %s, %d", keysFileName, version) + } + + return nil +} + +func (w *alphaWriteLog) DeleteTableKeys(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, version int) error { + bucket := partitionBucket(qtid.Key(), partition) + return w.wl.DeleteLog(bucket, keysFileName, version) +} + +func (w *alphaWriteLog) CreateFieldKeys(ctx context.Context, qtid dax.QualifiedTableID, field dax.FieldName, version int, m map[string]uint64) error { + msg := computer.FieldKeyMap{ + TableKey: qtid.Key(), + Field: field, + StringToID: m, + } + + b, err := json.Marshal(msg) + if err != nil { + return errors.Wrap(err, "marshalling field key map to json") + } + + bucket := fieldBucket(qtid.Key(), field) + + if err := w.wl.AppendMessage(bucket, keysFileName, version, b); err != nil { + return errors.Wrapf(err, "appending field key message: %s, %d", keysFileName, version) + } + + return nil +} + +func (w *alphaWriteLog) DeleteFieldKeys(ctx context.Context, qtid dax.QualifiedTableID, field dax.FieldName, version int) error { + bucket := fieldBucket(qtid.Key(), field) + return w.wl.DeleteLog(bucket, keysFileName, version) +} + +func (w *alphaWriteLog) WriteShard(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, shard dax.ShardNum, version int, msg computer.LogMessage) error { + b, err := computer.MarshalLogMessage(msg) + if err != nil { + return errors.Wrap(err, "marshalling log message") + } + + bucket := partitionBucket(qtid.Key(), partition) + shardKey := shardKey(shard) + + if err := w.wl.AppendMessage(bucket, shardKey, version, b); err != nil { + return errors.Wrapf(err, "appending shard key message: %s, %d", shardKey, version) + } + + return nil +} + +func (w *alphaWriteLog) DeleteShard(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, shard dax.ShardNum, version int) error { + bucket := partitionBucket(qtid.Key(), partition) + shardKey := shardKey(shard) + + return w.wl.DeleteLog(bucket, shardKey, version) +} + +//////////////////////////////////////////////// + +func (w *alphaWriteLog) TableKeyReader(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, version int) computer.TableKeyReader { + return newTableKeyReader(w.wl, qtid, partition, version) +} + +type tableKeyReader struct { + wl featurebase.WriteLogger + table dax.TableKey + partition dax.PartitionNum + version int + scanner *bufio.Scanner + closer io.Closer +} + +func newTableKeyReader(wl featurebase.WriteLogger, qtid dax.QualifiedTableID, partition dax.PartitionNum, version int) *tableKeyReader { + r := &tableKeyReader{ + wl: wl, + table: qtid.Key(), + partition: partition, + version: version, + } + + return r +} + +func (r *tableKeyReader) Open() error { + bucket := partitionBucket(r.table, r.partition) + + reader, closer, err := r.wl.LogReader(bucket, keysFileName, r.version) + if err != nil { + return errors.Wrapf(err, "getting log reader: %s, %s, %d", bucket, keysFileName, r.version) + } + + r.closer = closer + r.scanner = bufio.NewScanner(reader) + + return nil +} + +func (r *tableKeyReader) Read() (computer.PartitionKeyMap, error) { + if r.scanner == nil { + return computer.PartitionKeyMap{}, io.EOF + } + + var b []byte + var out computer.PartitionKeyMap + + if r.scanner.Scan() { + b = r.scanner.Bytes() + if err := json.Unmarshal(b, &out); err != nil { + return out, err + } + return out, nil + } + if err := r.scanner.Err(); err != nil { + return out, err + } + + return out, io.EOF +} + +func (r *tableKeyReader) Close() error { + if r.closer != nil { + return r.closer.Close() + } + return nil +} + +//////////////////////////////////////////////// + +func (w *alphaWriteLog) FieldKeyReader(ctx context.Context, qtid dax.QualifiedTableID, field dax.FieldName, version int) computer.FieldKeyReader { + return newFieldKeyReader(w.wl, qtid, field, version) +} + +type fieldKeyReader struct { + wl featurebase.WriteLogger + table dax.TableKey + field dax.FieldName + version int + scanner *bufio.Scanner + closer io.Closer +} + +func newFieldKeyReader(wl featurebase.WriteLogger, qtid dax.QualifiedTableID, field dax.FieldName, version int) *fieldKeyReader { + r := &fieldKeyReader{ + wl: wl, + table: qtid.Key(), + field: field, + version: version, + } + + return r +} + +func (r *fieldKeyReader) Open() error { + bucket := fieldBucket(r.table, r.field) + + reader, closer, err := r.wl.LogReader(bucket, keysFileName, r.version) + if err != nil { + return errors.Wrapf(err, "getting log reader: %s, %s, %d", bucket, keysFileName, r.version) + } + + r.closer = closer + r.scanner = bufio.NewScanner(reader) + + return nil +} + +func (r *fieldKeyReader) Read() (computer.FieldKeyMap, error) { + if r.scanner == nil { + return computer.FieldKeyMap{}, io.EOF + } + + var b []byte + var out computer.FieldKeyMap + + if r.scanner.Scan() { + b = r.scanner.Bytes() + if err := json.Unmarshal(b, &out); err != nil { + return out, err + } + return out, nil + } + if err := r.scanner.Err(); err != nil { + return out, err + } + + return out, io.EOF +} + +func (r *fieldKeyReader) Close() error { + if r.closer != nil { + return r.closer.Close() + } + return nil +} + +//////////////////////////////////////////////// + +func (w *alphaWriteLog) ShardReader(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, shard dax.ShardNum, version int) computer.ShardReader { + return newShardReader(w.wl, qtid, partition, shard, version) +} + +type shardReader struct { + wl featurebase.WriteLogger + table dax.TableKey + partition dax.PartitionNum + shard dax.ShardNum + version int + scanner *bufio.Scanner + closer io.Closer +} + +func newShardReader(wl featurebase.WriteLogger, qtid dax.QualifiedTableID, partition dax.PartitionNum, shard dax.ShardNum, version int) *shardReader { + r := &shardReader{ + wl: wl, + table: qtid.Key(), + partition: partition, + shard: shard, + version: version, + } + + return r +} + +func (r *shardReader) Open() error { + bucket := partitionBucket(r.table, r.partition) + shardKey := shardKey(r.shard) + + reader, closer, err := r.wl.LogReader(bucket, shardKey, r.version) + if err != nil { + return errors.Wrapf(err, "getting log reader: %s, %s, %d", bucket, shardKey, r.version) + } + + r.closer = closer + r.scanner = bufio.NewScanner(reader) + + return nil +} + +func (r *shardReader) Read() (computer.LogMessage, error) { + if r.scanner == nil { + return nil, io.EOF + } + + if r.scanner.Scan() { + b := r.scanner.Bytes() + + if len(b) == 0 { + return nil, errors.New(errors.ErrUncoded, "empty log record") + } + logMessageType := b[0] + + msg, err := computer.LogMessageByType(logMessageType) + if err != nil { + return nil, errors.Wrap(err, "getting log message by type") + } + + if err := json.Unmarshal(b[1:], &msg); err != nil { + return nil, errors.Wrap(err, "unmarshaling log message") + } + return msg, nil + } + if err := r.scanner.Err(); err != nil { + return nil, err + } + + return nil, io.EOF +} + +func (r *shardReader) Close() error { + if r.closer != nil { + return r.closer.Close() + } + return nil +} diff --git a/dax/computer/api/openapi.yaml b/dax/computer/api/openapi.yaml new file mode 100644 index 000000000..b24b6c44e --- /dev/null +++ b/dax/computer/api/openapi.yaml @@ -0,0 +1,235 @@ +openapi: 3.0.3 + +info: + title: Computer + description: The dax-related API for the Computer service. + version: 0.0.0 + +paths: + /computer/health: + get: + summary: Health check endpoint. + description: Provides an endpoint to check the overall health of the Computer service. + operationId: GetHealth + responses: + 200: + description: Service is healthy. + + /computer/directive: + post: + summary: Post Directive to compute node. + description: Post a Directive to the compute node. + operationId: PostDirective + requestBody: + content: + application/json: + schema: + $ref: '#/components/schemas/Directive' + responses: + 200: + description: Directive was applied successfully. + + /computer/snapshot/shard-data: + post: + summary: Request to snapshot shard data. + description: Request to snapshot shard data. + operationId: PostSnapshotShardData + requestBody: + content: + application/json: + schema: + type: object + properties: + address: + type: string + table: + type: string + shard: + type: integer + format: int64 + fromVersion: + type: integer + format: int64 + toVersion: + type: integer + format: int64 + directive: + $ref: '#/components/schemas/Directive' + + responses: + 200: + description: Shard snapshot was successful. + + /computer/snapshot/table-keys: + post: + summary: Request to snapshot table keys. + description: Request to snapshot table keys. + operationId: PostSnapshotTableKeys + requestBody: + content: + application/json: + schema: + type: object + properties: + address: + type: string + table: + type: string + partition: + type: integer + format: int32 + fromVersion: + type: integer + format: int64 + toVersion: + type: integer + format: int64 + directive: + $ref: '#/components/schemas/Directive' + + responses: + 200: + description: Table keys snapshot was successful. + + /computer/snapshot/field-keys: + post: + summary: Request to snapshot field keys. + description: Request to snapshot field keys. + operationId: PostSnapshotFieldKeys + requestBody: + content: + application/json: + schema: + type: object + properties: + address: + type: string + table: + type: string + field: + type: string + fromVersion: + type: integer + format: int64 + toVersion: + type: integer + format: int64 + directive: + $ref: '#/components/schemas/Directive' + + responses: + 200: + description: Field keys snapshot was successful. + +components: + responses: + Directive: + description: Directive response. + content: + application/json: + schema: + $ref: '#/components/schemas/Directive' + + schemas: + Directive: + type: object + properties: + address: + type: string + tables: + type: array + items: + $ref: '#/components/schemas/Table' + computeRoles: + type: array + items: + type: object + properties: + table: + type: string + shards: + type: array + items: + type: integer + format: int64 + translateRoles: + type: array + items: + type: object + properties: + table: + type: string + partitions: + type: array + items: + type: integer + format: int32 + fields: + type: array + items: + type: string + version: + type: integer + format: int64 + + # This is copied from /mds/api/openapi.yaml. TODO: share schemas across yaml files. + Table: + type: object + properties: + name: + type: string + fields: + type: array + items: + $ref: '#/components/schemas/Field' + partitionN: + type: integer + format: int32 + + Field: + type: object + properties: + name: + type: string + type: + type: string + enum: + - bool + - decimal + - id + - idset + - int + - string + - stringset + - timestamp + options: + type: object + properties: + min: + type: integer + format: int64 + max: + type: integer + format: int64 + scale: + type: integer + format: int64 + minimum: 0 + noStandardView: + type: boolean + cacheType: + type: string + cacheSize: + type: integer + format: int32 + timeUnit: + type: string + epoch: + type: string + format: date-time + timeQuantum: + type: string + ttl: + type: string + foreignIndex: + type: string \ No newline at end of file diff --git a/dax/computer/computer.go b/dax/computer/computer.go new file mode 100644 index 000000000..304a80c38 --- /dev/null +++ b/dax/computer/computer.go @@ -0,0 +1,5 @@ +// Package computer contains the compute-specific portions of the DAX +// architecture. In general, this is a dumb FeatureBase node (or service) which +// essentially contains the Executor and its interaction with the underlying +// data. +package computer diff --git a/dax/computer/snapshot.go b/dax/computer/snapshot.go new file mode 100644 index 000000000..3b19e93ed --- /dev/null +++ b/dax/computer/snapshot.go @@ -0,0 +1,61 @@ +package computer + +import ( + "context" + "io" + + "github.com/molecula/featurebase/v3/dax" +) + +// SnapshotReadWriter provides the interface for all snapshot read and writes in +// FeatureBase. +type SnapshotReadWriter interface { + WriteShardData(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, shard dax.ShardNum, version int, rc io.ReadCloser) error + ReadShardData(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, shard dax.ShardNum, version int) (io.ReadCloser, error) + + WriteTableKeys(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, version int, wrTo io.WriterTo) error + ReadTableKeys(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, version int) (io.ReadCloser, error) + + WriteFieldKeys(ctx context.Context, qtid dax.QualifiedTableID, field dax.FieldName, version int, wrTo io.WriterTo) error + ReadFieldKeys(ctx context.Context, qtid dax.QualifiedTableID, field dax.FieldName, version int) (io.ReadCloser, error) +} + +// Ensure type implements interface. +var _ SnapshotReadWriter = &NopSnapshotReadWriter{} + +// NopSnapshotReadWriter is a no-op implementation of the SnapshotReadWriter +// interface. +type NopSnapshotReadWriter struct{} + +func NewNopSnapshotReadWriter() *NopSnapshotReadWriter { + return &NopSnapshotReadWriter{} +} + +func (w *NopSnapshotReadWriter) WriteShardData(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, shard dax.ShardNum, version int, rc io.ReadCloser) error { + return nil +} + +func (w *NopSnapshotReadWriter) ReadShardData(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, shard dax.ShardNum, version int) (io.ReadCloser, error) { + return &nopReadCloser{}, nil +} + +func (w *NopSnapshotReadWriter) WriteTableKeys(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, version int, wrTo io.WriterTo) error { + return nil +} + +func (w *NopSnapshotReadWriter) ReadTableKeys(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, version int) (io.ReadCloser, error) { + return &nopReadCloser{}, nil +} + +func (w *NopSnapshotReadWriter) WriteFieldKeys(ctx context.Context, qtid dax.QualifiedTableID, field dax.FieldName, version int, wrTo io.WriterTo) error { + return nil +} + +func (w *NopSnapshotReadWriter) ReadFieldKeys(ctx context.Context, qtid dax.QualifiedTableID, field dax.FieldName, version int) (io.ReadCloser, error) { + return &nopReadCloser{}, nil +} + +type nopReadCloser struct{} + +func (n *nopReadCloser) Read([]byte) (int, error) { return 0, nil } +func (n *nopReadCloser) Close() error { return nil } diff --git a/dax/computer/writelog.go b/dax/computer/writelog.go new file mode 100644 index 000000000..07706df3a --- /dev/null +++ b/dax/computer/writelog.go @@ -0,0 +1,307 @@ +package computer + +import ( + "context" + "encoding/json" + "io" + "time" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/errors" +) + +// WriteLogWriter provides the interface for all data writes to FeatureBase. After +// data has been written to the local FeatureBase node, the respective interface +// method(s) will be called. +type WriteLogWriter interface { + // CreateTableKeys sends a map of string key to uint64 ID for the table and + // partition provided. + CreateTableKeys(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, version int, _ map[string]uint64) error + DeleteTableKeys(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, version int) error + + // CreateFieldKeys sends a map of string key to uint64 ID for the table and + // field provided. + CreateFieldKeys(ctx context.Context, qtid dax.QualifiedTableID, field dax.FieldName, version int, _ map[string]uint64) error + DeleteFieldKeys(ctx context.Context, qtid dax.QualifiedTableID, field dax.FieldName, version int) error + + WriteShard(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, shard dax.ShardNum, version int, msg LogMessage) error + DeleteShard(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, shard dax.ShardNum, version int) error +} + +// Ensure type implements interface. +var _ WriteLogWriter = (*NopWriteLogWriter)(nil) + +// NopWriteLogWriter is a no-op implementation of the WriteLogWriter interface. +type NopWriteLogWriter struct{} + +func NewNopWriteLogWriter() *NopWriteLogWriter { + return &NopWriteLogWriter{} +} + +func (w *NopWriteLogWriter) CreateTableKeys(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, version int, m map[string]uint64) error { + return nil +} + +func (w *NopWriteLogWriter) DeleteTableKeys(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, version int) error { + return nil +} + +func (w *NopWriteLogWriter) CreateFieldKeys(ctx context.Context, qtid dax.QualifiedTableID, field dax.FieldName, version int, m map[string]uint64) error { + return nil +} + +func (w *NopWriteLogWriter) DeleteFieldKeys(ctx context.Context, qtid dax.QualifiedTableID, field dax.FieldName, version int) error { + return nil +} + +func (w *NopWriteLogWriter) WriteShard(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, shard dax.ShardNum, version int, msg LogMessage) error { + return nil +} + +func (w *NopWriteLogWriter) DeleteShard(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, shard dax.ShardNum, version int) error { + return nil +} + +// WriteLogReader provides the interface for all reads from the write log. +type WriteLogReader interface { + ShardReader(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, shard dax.ShardNum, version int) ShardReader + TableKeyReader(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, version int) TableKeyReader + FieldKeyReader(ctx context.Context, qtid dax.QualifiedTableID, field dax.FieldName, version int) FieldKeyReader +} + +// Ensure type implements interface. +var _ WriteLogReader = (*NopWriteLogReader)(nil) + +// NopWriteLogReader is a no-op implementation of the WriteLogReader interface. +type NopWriteLogReader struct{} + +func NewNopWriteLogReader() *NopWriteLogReader { + return &NopWriteLogReader{} +} + +func (w *NopWriteLogReader) TableKeyReader(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, version int) TableKeyReader { + return NewNopTableKeyReader() +} + +func (w *NopWriteLogReader) FieldKeyReader(ctx context.Context, qtid dax.QualifiedTableID, field dax.FieldName, version int) FieldKeyReader { + return NewNopFieldKeyReader() +} + +func (w *NopWriteLogReader) ShardReader(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum, shard dax.ShardNum, version int) ShardReader { + return NewNopShardReader() +} + +//////////////////////////////////////////////// + +type TableKeyReader interface { + Open() error + Read() (PartitionKeyMap, error) + Close() error +} + +// Ensure type implements interface. +var _ TableKeyReader = &NopTableKeyReader{} + +// NopTableKeyReader is a no-op implementation of the TableKeyReader +// interface. +type NopTableKeyReader struct{} + +func NewNopTableKeyReader() *NopTableKeyReader { + return &NopTableKeyReader{} +} + +func (r *NopTableKeyReader) Open() error { return nil } +func (r *NopTableKeyReader) Read() (PartitionKeyMap, error) { + return PartitionKeyMap{}, io.EOF +} +func (r *NopTableKeyReader) Close() error { return nil } + +//////////////////////////////////////////////// + +type FieldKeyReader interface { + Open() error + Read() (FieldKeyMap, error) + Close() error +} + +// Ensure type implements interface. +var _ FieldKeyReader = &NopFieldKeyReader{} + +// NopFieldKeyReader is a no-op implementation of the FieldKeyReader +// interface. +type NopFieldKeyReader struct{} + +func NewNopFieldKeyReader() *NopFieldKeyReader { + return &NopFieldKeyReader{} +} + +func (r *NopFieldKeyReader) Open() error { return nil } +func (r *NopFieldKeyReader) Read() (FieldKeyMap, error) { + return FieldKeyMap{}, io.EOF +} +func (r *NopFieldKeyReader) Close() error { return nil } + +//////////////////////////////////////////////// + +type ShardReader interface { + Open() error + Read() (LogMessage, error) + Close() error +} + +// Ensure type implements interface. +var _ ShardReader = &NopShardReader{} + +// NopShardReader is a no-op implementation of the ShardReader interface. +type NopShardReader struct{} + +func NewNopShardReader() *NopShardReader { + return &NopShardReader{} +} + +func (r *NopShardReader) Open() error { return nil } +func (r *NopShardReader) Read() (LogMessage, error) { + return nil, io.EOF +} +func (r *NopShardReader) Close() error { return nil } + +//////////////// Messages /////////////////////// + +type PartitionKeyMap struct { + TableKey dax.TableKey `json:"table-key"` + Partition dax.PartitionNum `json:"partition"` + StringToID map[string]uint64 `json:"string-to-id"` +} + +type FieldKeyMap struct { + TableKey dax.TableKey `json:"table-key"` + Field dax.FieldName `json:"field"` + StringToID map[string]uint64 `json:"string-to-id"` +} + +const ( + logMessageTypeImportRoaring = iota + logMessageTypeImport + logMessageTypeImportValue + logMessageTypeImportRoaringShard +) + +type LogMessage interface{} + +// MarshalLogMessage serializes the log message and adds log message type info. +func MarshalLogMessage(msg LogMessage) ([]byte, error) { + typ, err := getLogMessageType(msg) + if err != nil { + return nil, errors.Wrap(err, "getting log message type") + } + + buf, err := json.Marshal(msg) + if err != nil { + return nil, errors.Wrap(err, "marshaling log message") + } + return append([]byte{typ}, buf...), nil +} + +func LogMessageByType(typ byte) (LogMessage, error) { + switch typ { + case logMessageTypeImportRoaring: + return &ImportRoaringMessage{}, nil + case logMessageTypeImport: + return &ImportMessage{}, nil + case logMessageTypeImportValue: + return &ImportValueMessage{}, nil + case logMessageTypeImportRoaringShard: + return &ImportRoaringShardMessage{}, nil + default: + return nil, errors.Errorf("unknown message type %d", typ) + } +} + +func getLogMessageType(m LogMessage) (byte, error) { + switch m.(type) { + case *ImportRoaringMessage: + return logMessageTypeImportRoaring, nil + case *ImportMessage: + return logMessageTypeImport, nil + case *ImportValueMessage: + return logMessageTypeImportValue, nil + case *ImportRoaringShardMessage: + return logMessageTypeImportRoaringShard, nil + default: + return 0, errors.Errorf("don't have type for message %#v", m) + } +} + +type ImportRoaringMessage struct { + LogMessage `json:"-"` + + Table string `json:"table"` + Field string `json:"field"` + Partition int `json:"partition"` + Shard uint64 `json:"shard"` + Clear bool `json:"clear"` + Action string `json:"action"` // [set, clear, overwrite] + Block int `json:"block"` + Views map[string][]byte `json:"views"` + UpdateExistence bool `json:"update-existence"` +} + +type ImportMessage struct { + LogMessage `json:"-"` + + Table string `json:"table"` + Field string `json:"field"` + Partition int `json:"partition"` + Shard uint64 `json:"shard"` + RowIDs []uint64 `json:"row-ids"` + ColumnIDs []uint64 `json:"column-ids"` + RowKeys []string `json:"row-keys"` + ColumnKeys []string `json:"column-keys"` + Timestamps []int64 `json:"timestamps"` + Clear bool `json:"clear"` + + // options + IgnoreKeyCheck bool `json:"ignore-key-check"` + Presorted bool `json:"presorted"` +} + +type ImportValueMessage struct { + LogMessage `json:"-"` + + Table string `json:"table"` + Field string `json:"field"` + Partition int `json:"partition"` + Shard uint64 `json:"shard"` + ColumnIDs []uint64 `json:"column-ids"` + ColumnKeys []string `json:"column-keys"` + Values []int64 `json:"values"` + FloatValues []float64 `json:"float-values"` + TimestampValues []time.Time `json:"timestamp-values"` + StringValues []string `json:"string-values"` + Clear bool `json:"clear"` + + // options + IgnoreKeyCheck bool `json:"ignore-key-check"` + Presorted bool `json:"presorted"` +} + +type ImportRoaringShardMessage struct { + LogMessage `json:"-"` + + Table string `json:"table"` + Partition int `json:"partition"` + Shard uint64 `json:"shard"` + Views []RoaringUpdate `json:"views"` +} + +// RoaringUpdate is identical to featurebase.RoaringUpdate, but we +// can't import it due to import cycles. TODO featurebase top level +// shouldn't import dax stuff... all the types it needs should just be +// in the top level. +type RoaringUpdate struct { + Field string `json:"field"` + View string `json:"view"` + Clear []byte `json:"clear"` + Set []byte `json:"set"` + ClearRecords bool `json:"clear-records"` +} diff --git a/dax/dax.go b/dax/dax.go new file mode 100644 index 000000000..052969d87 --- /dev/null +++ b/dax/dax.go @@ -0,0 +1,11 @@ +// Package dax defines DAX domain level types. +package dax + +// ServicePrefixes are used as the service prefix value in http handlers. +const ( + ServicePrefixComputer = "computer" + ServicePrefixMDS = "mds" + ServicePrefixQueryer = "queryer" + ServicePrefixSnapshotter = "snapshotter" + ServicePrefixWriteLogger = "writelogger" +) diff --git a/dax/directive.go b/dax/directive.go new file mode 100644 index 000000000..869bbdade --- /dev/null +++ b/dax/directive.go @@ -0,0 +1,175 @@ +package dax + +// Directive contains the instructions, sent from MDS, which a compute node is +// to follow. A Directive is typically JSON-encoded and POSTed to a compute +// node's `/directive` endpoint. +type Directive struct { + Address Address `json:"address"` + + // Method describes how the compute node should handle the Directive. See + // the different constants of type DirectiveMethod for how this value is + // handled. + Method DirectiveMethod `json:"method"` + + Tables []*QualifiedTable `json:"schema"` + + ComputeRoles []ComputeRole `json:"compute-roles"` + TranslateRoles []TranslateRole `json:"translate-roles"` + + Version uint64 `json:"version"` +} + +// DirectiveMethod is used to tell the compute node how it should handle the +// Directive. +type DirectiveMethod string + +const ( + // DirectiveMethodDiff tells the compute node to diff the Directive with its + // local, cached Directive and only apply the differences. + DirectiveMethodDiff DirectiveMethod = "diff" + + // DirectiveMethodReset tells the compute node to delete all of its existing + // data before applying the directive. + DirectiveMethodReset DirectiveMethod = "reset" + + // DirectiveMethodSnapshot tells the compute node that the incoming + // Directive should only contain data version updates related to a snapshot + // request. + DirectiveMethodSnapshot DirectiveMethod = "snapshot" +) + +// Table returns the ID'd table from the Directive's Tables list. If it's not +// found, it returns nil and a non-nil error. A nil error guarantees that the +// returned table is non-nil. +func (d *Directive) Table(qtid QualifiedTableID) (*QualifiedTable, error) { + for _, qtbl := range d.Tables { + // We can't do qtbl.QualifiedID() == qtid because the value of qtid.Name + // is empty and causes the equality check to fail. Hence the .Equals() + // method. + if qtbl.QualifiedID().Equals(qtid) { + return qtbl, nil + } + } + return nil, NewErrTableIDDoesNotExist(qtid) +} + +// ComputeShards returns the list of shards, for the given table, for which this +// compute node is responsible. It assumes that the Directive does not contain +// more than one ComputeRole for the same table; in that case, we would need to +// return the union of Shards. +func (d *Directive) ComputeShards(tbl TableKey) Shards { + if d == nil || d.ComputeRoles == nil { + return Shards{} + } + + for _, cr := range d.ComputeRoles { + if cr.TableKey == tbl { + return cr.Shards + } + } + + return Shards{} +} + +// ComputeShardsMap returns a map of table to shards. It assumes that the +// Directive does not contain more than one ComputeRole for the same table; in +// that case, we would need to return the union of Shards. +func (d *Directive) ComputeShardsMap() map[TableKey]Shards { + m := make(map[TableKey]Shards) + if d == nil || d.ComputeRoles == nil { + return m + } + + for _, cr := range d.ComputeRoles { + m[cr.TableKey] = cr.Shards + } + + return m +} + +// TranslatePartitions returns the list of partitions, for the given table, for +// which this translate node is responsible. It assumes that the Directive does +// not contain more than one TranslateRole for the same table; in that case, we +// would need to return the union of Shards. +func (d *Directive) TranslatePartitions(tbl TableKey) Partitions { + if d == nil || d.TranslateRoles == nil { + return Partitions{} + } + + for _, tr := range d.TranslateRoles { + if tr.TableKey == tbl { + return tr.Partitions + } + } + return Partitions{} +} + +// TranslatePartitionsMap returns a map of table to partitions. It assumes that +// the Directive does not contain more than one TranslateRole for the same +// table; in that case, we would need to return the union of Partitions. +func (d *Directive) TranslatePartitionsMap() map[TableKey]Partitions { + m := make(map[TableKey]Partitions) + if d == nil || d.TranslateRoles == nil { + return m + } + + for _, tr := range d.TranslateRoles { + // Since we added FieldVersions to the TranslateRole, it's possible for + // a TranslateRole to have an empty Partitions list. In that case, we + // want to exclude that from the map. + if len(tr.Partitions) == 0 { + continue + } + m[tr.TableKey] = tr.Partitions + } + + return m +} + +// TranslateFieldsMap returns a map of table to fields. It assumes that +// the Directive does not contain more than one TranslateRole for the same +// table; in that case, we would need to return the union of FieldValues. +func (d *Directive) TranslateFieldsMap() map[TableKey]FieldVersions { + m := make(map[TableKey]FieldVersions) + if d == nil || d.TranslateRoles == nil { + return m + } + + for _, tr := range d.TranslateRoles { + if len(tr.Fields) == 0 { + continue + } + m[tr.TableKey] = tr.Fields + } + + return m +} + +// IsEmpty tells whether a directive is assigning actual responsibilty +// to a node or not. If the directive does not assign responsibility +// for any shard or partition then it is considered empty. This is +// used to determine whether we can ignore an error received from +// applying this directive (an empty directive is often sent to a node +// which is already down). +func (d *Directive) IsEmpty() bool { + for _, role := range d.ComputeRoles { + if len(role.Shards) > 0 { + return false + } + } + + for _, role := range d.TranslateRoles { + if len(role.Partitions) > 0 { + return false + } + } + + return true +} + +// Directives is a sortable slice of Directive. +type Directives []*Directive + +func (d Directives) Len() int { return len(d) } +func (d Directives) Less(i, j int) bool { return d[i].Address.String() < d[j].Address.String() } +func (d Directives) Swap(i, j int) { d[i], d[j] = d[j], d[i] } diff --git a/dax/docker-compose.yml b/dax/docker-compose.yml new file mode 100644 index 000000000..69e995297 --- /dev/null +++ b/dax/docker-compose.yml @@ -0,0 +1,66 @@ +version: '3' + +services: + mds: + build: + context: ../.quick + dockerfile: ../Dockerfile-dax-quick + environment: + FEATUREBASE_BIND: 0.0.0.0:8080 + FEATUREBASE_VERBOSE: "true" + FEATUREBASE_STORAGE_METHOD: boltdb + FEATUREBASE_STORAGE_DSN: file:/dax-data/mds.boldtb + FEATUREBASE_MDS_RUN: "true" + ports: + - "8081:8080" + + queryer: + build: + context: ../.quick + dockerfile: ../Dockerfile-dax-quick + environment: + FEATUREBASE_BIND: 0.0.0.0:8080 + FEATUREBASE_VERBOSE: "true" + FEATUREBASE_QUERYER_RUN: "true" + FEATUREBASE_QUERYER_CONFIG_MDS_ADDRESS: "mds:8080" + depends_on: + - mds + ports: + - "8080:8080" + + computer: + build: + context: ../.quick + dockerfile: ../Dockerfile-dax-quick + environment: + FEATUREBASE_COMPUTER_RUN: "true" + FEATUREBASE_COMPUTER_CONFIG_MDS_ADDRESS: "mds:8080" + FEATUREBASE_COMPUTER_CONFIG_DATA_DIR: /dax-data/computer + FEATUREBASE_BIND: 0.0.0.0:8080 + FEATUREBASE_VERBOSE: "true" + FEATUREBASE_STORAGE_METHOD: boltdb + FEATUREBASE_WRITELOGGER_RUN: "true" + FEATUREBASE_WRITELOGGER_CONFIG_DATA_DIR: "/dax-data/writelogger" + FEATUREBASE_SNAPSHOTTER_RUN: "true" + FEATUREBASE_SNAPSHOTTER_CONFIG_DATA_DIR: "/dax-data/snapshotter" + volumes: + - "./dax-data/writelogger:/dax-data/writelogger" + - "./dax-data/snapshotter:/dax-data/snapshotter" + depends_on: + - mds + deploy: + replicas: 1 + + datagen: + build: + context: .. + dockerfile: Dockerfile-datagen + profiles: [ "datagen" ] + environment: + GEN_CUSTOM_CONFIG: "/testdata/keys_ids.yaml" + GEN_FEATUREBASE_ORG_ID: "testorg" + GEN_FEATUREBASE_DB_ID: "testdb" + GEN_USE_SHARD_TRANSACTIONAL_ENDPOINT: "true" + GEN_SOURCE: "custom" + GEN_TARGET: "mds" + GEN_MDS_ADDRESS: "mds:8080" diff --git a/dax/errors.go b/dax/errors.go new file mode 100644 index 000000000..776fe5df9 --- /dev/null +++ b/dax/errors.go @@ -0,0 +1,80 @@ +package dax + +import ( + "fmt" + + "github.com/molecula/featurebase/v3/errors" +) + +const ( + ErrTableIDExists errors.Code = "TableIDExists" + ErrTableKeyExists errors.Code = "TableKeyExists" + ErrTableNameExists errors.Code = "TableNameExists" + ErrTableIDDoesNotExist errors.Code = "TableIDDoesNotExist" + ErrTableKeyDoesNotExist errors.Code = "TableKeyDoesNotExist" + ErrTableNameDoesNotExist errors.Code = "TableNameDoesNotExist" + + ErrFieldExists errors.Code = "FieldExists" + ErrFieldDoesNotExist errors.Code = "FieldDoesNotExist" + + ErrUnimplemented errors.Code = "Unimplemented" +) + +// The following are helper functions for constructing coded errors containing +// relevant information about the specific error. + +func NewErrTableIDDoesNotExist(qtid QualifiedTableID) error { + return errors.New( + ErrTableIDDoesNotExist, + fmt.Sprintf("table ID '%s' does not exist", qtid), + ) +} + +func NewErrTableKeyDoesNotExist(tkey TableKey) error { + return errors.New( + ErrTableKeyDoesNotExist, + fmt.Sprintf("table key '%s' does not exist", tkey), + ) +} + +func NewErrTableNameDoesNotExist(tableName TableName) error { + return errors.New( + ErrTableNameDoesNotExist, + fmt.Sprintf("table name '%s' does not exist", tableName), + ) +} + +func NewErrTableIDExists(qtid QualifiedTableID) error { + return errors.New( + ErrTableIDExists, + fmt.Sprintf("table ID '%s' already exists", qtid), + ) +} + +func NewErrTableKeyExists(tkey TableKey) error { + return errors.New( + ErrTableKeyExists, + fmt.Sprintf("table key '%s' already exists", tkey), + ) +} + +func NewErrTableNameExists(tableName TableName) error { + return errors.New( + ErrTableNameExists, + fmt.Sprintf("table name '%s' already exists", tableName), + ) +} + +func NewErrFieldDoesNotExist(fieldName FieldName) error { + return errors.New( + ErrFieldDoesNotExist, + fmt.Sprintf("field '%s' does not exist", fieldName), + ) +} + +func NewErrFieldExists(fieldName FieldName) error { + return errors.New( + ErrFieldExists, + fmt.Sprintf("field '%s' already exists", fieldName), + ) +} diff --git a/dax/fieldversion.go b/dax/fieldversion.go new file mode 100644 index 000000000..8f1133c57 --- /dev/null +++ b/dax/fieldversion.go @@ -0,0 +1,32 @@ +package dax + +import "fmt" + +// FieldVersion is used in a similar way to Shard and Partition in that they all +// contain a snapshot version. It would have been confusing to use the Field +// type which already exists, because versioning that would mean something else. +// This is really snapshot specific (as are Shard and Partition). +type FieldVersion struct { + Name FieldName `json:"name"` + Version int `json:"version"` +} + +// String returns the FieldVersion (i.e. its Name and Version) as a string. +func (f FieldVersion) String() string { + return fmt.Sprintf("%s.%d", f.Name, f.Version) +} + +// NewFieldVersion returns a FieldVersion with the provided name and version. +func NewFieldVersion(name FieldName, version int) FieldVersion { + return FieldVersion{ + Name: name, + Version: version, + } +} + +// FieldVersions is a sortable slice of FieldVersion. +type FieldVersions []FieldVersion + +func (f FieldVersions) Len() int { return len(f) } +func (f FieldVersions) Less(i, j int) bool { return f[i].Name < f[j].Name } +func (f FieldVersions) Swap(i, j int) { f[i], f[j] = f[j], f[i] } diff --git a/dax/http/handler.go b/dax/http/handler.go new file mode 100644 index 000000000..637201ac2 --- /dev/null +++ b/dax/http/handler.go @@ -0,0 +1,219 @@ +package http + +import ( + "context" + "net" + "net/http" + "runtime/debug" + "time" + + "github.com/gorilla/mux" + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/mds" + mdshttp "github.com/molecula/featurebase/v3/dax/mds/http" + "github.com/molecula/featurebase/v3/dax/queryer" + queryerhttp "github.com/molecula/featurebase/v3/dax/queryer/http" + "github.com/molecula/featurebase/v3/dax/snapshotter" + snapshotterhttp "github.com/molecula/featurebase/v3/dax/snapshotter/http" + "github.com/molecula/featurebase/v3/dax/writelogger" + writeloggerhttp "github.com/molecula/featurebase/v3/dax/writelogger/http" + "github.com/molecula/featurebase/v3/errors" + "github.com/molecula/featurebase/v3/logger" +) + +// Handler represents an HTTP handler. +type Handler struct { + Handler http.Handler + + bind string + + ln net.Listener + // url is used to hold the advertise bind address for printing a log during startup. + url string + + closeTimeout time.Duration + + server *http.Server + + mds *mds.MDS + writeLogger *writelogger.WriteLogger + snapshotter *snapshotter.Snapshotter + queryer *queryer.Queryer + + computer http.Handler + + logger logger.Logger +} + +// HandlerOption is a functional option type for Handler +type HandlerOption func(s *Handler) error + +func OptHandlerBind(b string) HandlerOption { + return func(h *Handler) error { + h.bind = b + return nil + } +} + +func OptHandlerMDS(m *mds.MDS) HandlerOption { + return func(h *Handler) error { + h.mds = m + return nil + } +} + +func OptHandlerWriteLogger(w *writelogger.WriteLogger) HandlerOption { + return func(h *Handler) error { + h.writeLogger = w + return nil + } +} + +func OptHandlerSnapshotter(s *snapshotter.Snapshotter) HandlerOption { + return func(h *Handler) error { + h.snapshotter = s + return nil + } +} + +func OptHandlerQueryer(q *queryer.Queryer) HandlerOption { + return func(h *Handler) error { + h.queryer = q + return nil + } +} + +func OptHandlerLogger(l logger.Logger) HandlerOption { + return func(h *Handler) error { + h.logger = l + return nil + } +} + +// OptHandlerCloseTimeout controls how long to wait for the http Server to +// shutdown cleanly before forcibly destroying it. Default is 30 seconds. +func OptHandlerCloseTimeout(d time.Duration) HandlerOption { + return func(h *Handler) error { + h.closeTimeout = d + return nil + } +} + +// OptHandlerListener set the listener that will be used by the HTTP server. +// Url must be the advertised URL. It will be used to show a log to the user +// about where the Web UI is. This option is mandatory. +func OptHandlerListener(ln net.Listener, url string) HandlerOption { + return func(h *Handler) error { + h.ln = ln + h.url = url + return nil + } +} + +func OptHandlerComputer(handler http.Handler) HandlerOption { + return func(h *Handler) error { + h.computer = handler + return nil + } +} + +// NewHandler returns a new instance of Handler with a default logger. +func NewHandler(opts ...HandlerOption) (*Handler, error) { + handler := &Handler{ + logger: logger.NopLogger, + closeTimeout: time.Second * 30, + } + + for _, opt := range opts { + err := opt(handler) + if err != nil { + return nil, errors.Wrap(err, "applying option") + } + } + + handler.Handler = newRouter(handler) + + handler.server = &http.Server{Handler: handler} + + return handler, nil +} + +func (h *Handler) Serve() error { + err := h.server.Serve(h.ln) + if err != nil && err.Error() != "http: Server closed" { + h.logger.Errorf("HTTP handler terminated with error: %s\n", err) + return errors.Wrap(err, "serve http") + } + return nil +} + +// Close tries to cleanly shutdown the HTTP server, and failing that, after a +// timeout, calls Server.Close. +func (h *Handler) Close() error { + deadlineCtx, cancelFunc := context.WithDeadline(context.Background(), time.Now().Add(h.closeTimeout)) + defer cancelFunc() + err := h.server.Shutdown(deadlineCtx) + if err != nil { + err = h.server.Close() + } + return errors.Wrap(err, "shutdown/close http server") +} + +// newRouter creates a new mux http router. +func newRouter(handler *Handler) http.Handler { + router := mux.NewRouter() + + router.HandleFunc("/health", handler.handleGetHealth).Methods("GET").Name("GetHealth") + + if handler.mds != nil { + pre := "/" + dax.ServicePrefixMDS + router.PathPrefix(pre).Handler( + http.StripPrefix(pre, mdshttp.Handler(handler.mds))) + } + + if handler.writeLogger != nil { + pre := "/" + dax.ServicePrefixWriteLogger + router.PathPrefix(pre).Handler( + http.StripPrefix(pre, writeloggerhttp.Handler(handler.writeLogger, handler.logger))) + } + + if handler.snapshotter != nil { + pre := "/" + dax.ServicePrefixSnapshotter + router.PathPrefix(pre).Handler( + http.StripPrefix(pre, snapshotterhttp.Handler(handler.snapshotter))) + } + + if handler.queryer != nil { + pre := "/" + dax.ServicePrefixQueryer + router.PathPrefix(pre).Handler( + http.StripPrefix(pre, queryerhttp.Handler(handler.queryer))) + } + + if handler.computer != nil { + pre := "/" + dax.ServicePrefixComputer + router.PathPrefix(pre).Handler( + http.StripPrefix(pre, handler.computer)) + } + + var h http.Handler = router + + return h +} + +// ServeHTTP handles an HTTP request. +func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { + defer func() { + if err := recover(); err != nil { + w.WriteHeader(http.StatusInternalServerError) + stack := debug.Stack() + h.logger.Printf("PANIC: %s\n%s", err, stack) + } + }() + + h.Handler.ServeHTTP(w, r) +} + +// GET /health +func (h *Handler) handleGetHealth(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) +} diff --git a/dax/inmem/inmem.go b/dax/inmem/inmem.go new file mode 100644 index 000000000..9a2663214 --- /dev/null +++ b/dax/inmem/inmem.go @@ -0,0 +1,2 @@ +// Package inmem contains the in-memory implementation of the dax interfaces. +package inmem diff --git a/dax/inmem/versionstore.go b/dax/inmem/versionstore.go new file mode 100644 index 000000000..5f30e41ce --- /dev/null +++ b/dax/inmem/versionstore.go @@ -0,0 +1,430 @@ +package inmem + +import ( + "context" + "sort" + "sync" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/errors" +) + +// Ensure type implements interface. +var _ dax.VersionStore = (*VersionStore)(nil) + +// VersionStore manages all version info for shard, table keys, and field keys. +type VersionStore struct { + mu sync.RWMutex + + // shards is a map of all shards, by table, by shard number, known to + // contain data. + shards map[dax.TableQualifierKey]map[dax.TableID]map[dax.ShardNum]dax.Shard + + // tableKeys is a map of all partitions, by table, by partition number, + // known to contain key data. + tableKeys map[dax.TableQualifierKey]map[dax.TableID]map[dax.PartitionNum]int + + // fieldKeys is a map of all fields, by table, known to contain key data. + fieldKeys map[dax.TableQualifierKey]map[dax.TableID]map[dax.FieldName]int +} + +// NewVersionStore returns a new instance of VersionStore with default values. +func NewVersionStore() *VersionStore { + return &VersionStore{ + shards: make(map[dax.TableQualifierKey]map[dax.TableID]map[dax.ShardNum]dax.Shard), + tableKeys: make(map[dax.TableQualifierKey]map[dax.TableID]map[dax.PartitionNum]int), + fieldKeys: make(map[dax.TableQualifierKey]map[dax.TableID]map[dax.FieldName]int), + } +} + +// AddTable adds a table to be managed by VersionStore. +func (s *VersionStore) AddTable(ctx context.Context, qtid dax.QualifiedTableID) error { + s.mu.Lock() + defer s.mu.Unlock() + + // This check is clunky; three maps contain the table, but we only check for + // existence in one of them. It also seems weird to check all three, because + // if we get in a state where one of the maps doesn't contain a table that + // the other maps do contain, the state of the data is in question. + if _, found := s.shards[qtid.TableQualifier.Key()][qtid.ID]; found { + return dax.NewErrTableIDExists(qtid) + } + + // Initialize the maps in case VersionStore wasn't created with NewVersionStore(). + if s.shards == nil { + s.shards = make(map[dax.TableQualifierKey]map[dax.TableID]map[dax.ShardNum]dax.Shard) + } + if s.tableKeys == nil { + s.tableKeys = make(map[dax.TableQualifierKey]map[dax.TableID]map[dax.PartitionNum]int) + } + if s.fieldKeys == nil { + s.fieldKeys = make(map[dax.TableQualifierKey]map[dax.TableID]map[dax.FieldName]int) + } + + // shards. + if _, ok := s.shards[qtid.TableQualifier.Key()]; !ok { + s.shards[qtid.TableQualifier.Key()] = make(map[dax.TableID]map[dax.ShardNum]dax.Shard, 0) + } + if _, ok := s.shards[qtid.TableQualifier.Key()][qtid.ID]; !ok { + s.shards[qtid.TableQualifier.Key()][qtid.ID] = make(map[dax.ShardNum]dax.Shard, 0) + } + + // tableKeys. + if _, ok := s.tableKeys[qtid.TableQualifier.Key()]; !ok { + s.tableKeys[qtid.TableQualifier.Key()] = make(map[dax.TableID]map[dax.PartitionNum]int, 0) + } + if _, ok := s.tableKeys[qtid.TableQualifier.Key()][qtid.ID]; !ok { + s.tableKeys[qtid.TableQualifier.Key()][qtid.ID] = make(map[dax.PartitionNum]int, 0) + } + + // fieldKeys. + if _, ok := s.fieldKeys[qtid.TableQualifier.Key()]; !ok { + s.fieldKeys[qtid.TableQualifier.Key()] = make(map[dax.TableID]map[dax.FieldName]int, 0) + } + if _, ok := s.fieldKeys[qtid.TableQualifier.Key()][qtid.ID]; !ok { + s.fieldKeys[qtid.TableQualifier.Key()][qtid.ID] = make(map[dax.FieldName]int, 0) + } + + return nil +} + +// RemoveTable removes the given table. An error will be returned if the table +// does not exist. +func (s *VersionStore) RemoveTable(ctx context.Context, qtid dax.QualifiedTableID) (dax.Shards, dax.Partitions, error) { + s.mu.Lock() + defer s.mu.Unlock() + + var foundTable bool + var shards dax.Shards + var partitions dax.Partitions + var err error + + // Remove shards for table. + if s.shards != nil { + if _, ok := s.shards[qtid.TableQualifier.Key()][qtid.ID]; ok { + foundTable = true + + // Get the shards to return before deleting from map. + shards, _, err = s.shardSlice(qtid) + if err != nil { + return nil, nil, errors.Wrapf(err, "getting shard slice: %s", qtid) + } + + // Remove the shards. + delete(s.shards[qtid.TableQualifier.Key()], qtid.ID) + } + } + + // Remove tableKeys for table. + if s.tableKeys != nil { + if _, ok := s.tableKeys[qtid.TableQualifier.Key()][qtid.ID]; ok { + foundTable = true + + // Get the partitions to return before deleting from map. + partitions, _, err = s.partitionSlice(qtid) + if err != nil { + return nil, nil, errors.Wrapf(err, "getting partition slice: %s", qtid) + } + + // Remove the tableKeys. + delete(s.tableKeys[qtid.TableQualifier.Key()], qtid.ID) + } + } + + // Remove fieldKeys for table. + if s.fieldKeys != nil { + if _, ok := s.fieldKeys[qtid.TableQualifier.Key()][qtid.ID]; ok { + foundTable = true + + // Remove the fieldKeys. + delete(s.fieldKeys[qtid.TableQualifier.Key()], qtid.ID) + } + } + + if !foundTable { + return nil, nil, dax.NewErrTableIDDoesNotExist(qtid) + } + + return shards, partitions, nil +} + +// AddShards adds new shards to be managed by VersionStore. It returns the +// number of shards added or an error. +func (s *VersionStore) AddShards(ctx context.Context, qtid dax.QualifiedTableID, shards ...dax.Shard) error { + s.mu.Lock() + defer s.mu.Unlock() + + sh, ok := s.shards[qtid.TableQualifier.Key()][qtid.ID] + if !ok { + return dax.NewErrTableIDDoesNotExist(qtid) + } + + var n int + + for _, shard := range shards { + if _, ok := sh[shard.Num]; !ok { + n++ // TODO: this isn't considering a shard that exists, but the version changes. + } + sh[shard.Num] = shard + } + + return nil +} + +// Shards returns the list of shards available for the give table. It returns +// false if the table does not exist. +func (s *VersionStore) Shards(ctx context.Context, qtid dax.QualifiedTableID) (dax.Shards, bool, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + return s.shardSlice(qtid) +} + +// shardSlice is an unprotected version of Shards(). +func (s *VersionStore) shardSlice(qtid dax.QualifiedTableID) (dax.Shards, bool, error) { + if s.shards == nil { + return nil, false, nil + } + + if shardNumMap, ok := s.shards[qtid.TableQualifier.Key()][qtid.ID]; ok { + rtn := make(dax.Shards, 0, len(shardNumMap)) + for _, shard := range shardNumMap { + rtn = append(rtn, shard) + } + sort.Sort(rtn) + return rtn, true, nil + } + + return nil, false, nil +} + +// ShardVersion return the current version for the given table/shardNum. +// If a version is not being tracked, it returns a bool value of false. +func (s *VersionStore) ShardVersion(ctx context.Context, qtid dax.QualifiedTableID, shardNum dax.ShardNum) (int, bool, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + t, ok := s.shards[qtid.TableQualifier.Key()][qtid.ID] + if !ok { + return -1, false, nil + } + + v, ok := t[shardNum] + if !ok { + return -1, false, nil + } + return v.Version, true, nil +} + +func (s *VersionStore) ShardTables(ctx context.Context, qual dax.TableQualifier) (dax.TableIDs, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + qual.Key() + + tableIDs := make(dax.TableIDs, 0, len(s.shards[qual.Key()])) + + for tableID := range s.shards[qual.Key()] { + tableIDs = append(tableIDs, tableID) + } + + return tableIDs, nil +} + +// AddPartitions adds new partitions to be managed by VersionStore. It returns +// the number of partitions added or an error. +func (s *VersionStore) AddPartitions(ctx context.Context, qtid dax.QualifiedTableID, partitions ...dax.Partition) error { + s.mu.Lock() + defer s.mu.Unlock() + + tk, ok := s.tableKeys[qtid.TableQualifier.Key()][qtid.ID] + if !ok { + return dax.NewErrTableIDDoesNotExist(qtid) + } + + for _, partition := range partitions { + tk[partition.Num] = partition.Version + } + + return nil +} + +// Partitions returns the list of partitions available for the give table. It +// returns false if the table does not exist. +func (s *VersionStore) Partitions(ctx context.Context, qtid dax.QualifiedTableID) (dax.Partitions, bool, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + return s.partitionSlice(qtid) +} + +// partitionSlice is an unprotected version of Partitions(). +func (s *VersionStore) partitionSlice(qtid dax.QualifiedTableID) (dax.Partitions, bool, error) { + if s.tableKeys == nil { + return nil, false, nil + } + + if partitionNumMap, ok := s.tableKeys[qtid.TableQualifier.Key()][qtid.ID]; ok { + rtn := make(dax.Partitions, 0, len(partitionNumMap)) + for partitionNum, version := range partitionNumMap { + rtn = append(rtn, dax.NewPartition(partitionNum, version)) + } + sort.Sort(rtn) + return rtn, true, nil + } + + return nil, false, nil +} + +// PartitionVersion return the current version for the given table/partitionNum. +// If a version is not being tracked, it returns a bool value of false. +func (s *VersionStore) PartitionVersion(ctx context.Context, qtid dax.QualifiedTableID, partitionNum dax.PartitionNum) (int, bool, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + t, ok := s.tableKeys[qtid.TableQualifier.Key()][qtid.ID] + if !ok { + return -1, false, nil + } + + v, ok := t[partitionNum] + if !ok { + return -1, false, nil + } + return v, true, nil +} + +func (s *VersionStore) PartitionTables(ctx context.Context, qual dax.TableQualifier) (dax.TableIDs, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + tableIDs := make(dax.TableIDs, 0, len(s.tableKeys[qual.Key()])) + + for tableName := range s.tableKeys[qual.Key()] { + tableIDs = append(tableIDs, tableName) + } + + return tableIDs, nil +} + +// AddFields adds new fields to be managed by VersionStore. It returns the +// number of fields added or an error. +func (s *VersionStore) AddFields(ctx context.Context, qtid dax.QualifiedTableID, fields ...dax.FieldVersion) error { + s.mu.Lock() + defer s.mu.Unlock() + + fk, ok := s.fieldKeys[qtid.TableQualifier.Key()][qtid.ID] + if !ok { + return dax.NewErrTableIDDoesNotExist(qtid) + } + + for _, field := range fields { + fk[field.Name] = field.Version + } + + return nil +} + +// Fields returns the list of fields available for the give table. It returns +// false if the table does not exist. +func (s *VersionStore) Fields(ctx context.Context, qtid dax.QualifiedTableID) (dax.FieldVersions, bool, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + return s.fieldSlice(qtid) +} + +// fieldSlice is an unprotected version of Fields(). +func (s *VersionStore) fieldSlice(qtid dax.QualifiedTableID) (dax.FieldVersions, bool, error) { + if s.fieldKeys == nil { + return nil, false, nil + } + + if fieldNameMap, ok := s.fieldKeys[qtid.TableQualifier.Key()][qtid.ID]; ok { + rtn := make(dax.FieldVersions, 0, len(fieldNameMap)) + for fieldName, version := range fieldNameMap { + rtn = append(rtn, dax.NewFieldVersion(fieldName, version)) + } + sort.Sort(rtn) + return rtn, true, nil + } + + return nil, false, nil +} + +// FieldVersion return the current version for the given table/field. +// If a version is not being tracked, it returns a bool value of false. +func (s *VersionStore) FieldVersion(ctx context.Context, qtid dax.QualifiedTableID, field dax.FieldName) (int, bool, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + t, ok := s.fieldKeys[qtid.TableQualifier.Key()][qtid.ID] + if !ok { + return -1, false, nil + } + + v, ok := t[field] + if !ok { + return -1, false, nil + } + return v, true, nil +} + +func (s *VersionStore) FieldTables(ctx context.Context, qual dax.TableQualifier) (dax.TableIDs, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + tableIDs := make(dax.TableIDs, 0, len(s.fieldKeys[qual.Key()])) + + for tableID := range s.fieldKeys[qual.Key()] { + tableIDs = append(tableIDs, tableID) + } + + return tableIDs, nil +} + +// Copy returns a new copy of VersionStore. +func (s *VersionStore) Copy(ctx context.Context) (dax.VersionStore, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + new := NewVersionStore() + + // shards. + for qkey, tableIDs := range s.shards { + for tableID, shards := range tableIDs { + qual := dax.NewTableQualifier(qkey.OrganizationID(), qkey.DatabaseID()) + qtid := dax.NewQualifiedTableID(qual, tableID) + _ = new.AddTable(ctx, qtid) + for shardNum, shard := range shards { + new.shards[qual.Key()][tableID][shardNum] = shard + } + } + } + + // tableKeys. + for qkey, tableIDs := range s.tableKeys { + for tableID, partitions := range tableIDs { + qual := dax.NewTableQualifier(qkey.OrganizationID(), qkey.DatabaseID()) + qtid := dax.NewQualifiedTableID(qual, tableID) + _ = new.AddTable(ctx, qtid) + for partitionNum, version := range partitions { + new.tableKeys[qual.Key()][tableID][partitionNum] = version + } + } + } + + // fieldKeys. + for qkey, tableIDs := range s.fieldKeys { + for tableID, fields := range tableIDs { + qual := dax.NewTableQualifier(qkey.OrganizationID(), qkey.DatabaseID()) + qtid := dax.NewQualifiedTableID(qual, tableID) + _ = new.AddTable(ctx, qtid) + for field, version := range fields { + new.fieldKeys[qual.Key()][tableID][field] = version + } + } + } + + return new, nil +} diff --git a/dax/inmem/versionstore_test.go b/dax/inmem/versionstore_test.go new file mode 100644 index 000000000..82db23b7b --- /dev/null +++ b/dax/inmem/versionstore_test.go @@ -0,0 +1,166 @@ +package inmem_test + +import ( + "context" + "testing" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/inmem" + "github.com/molecula/featurebase/v3/errors" + "github.com/stretchr/testify/assert" +) + +func TestVersionStore(t *testing.T) { + orgID := dax.OrganizationID("acme") + dbID := dax.DatabaseID("db1") + + tableID := dax.TableID("0000000000000001") + + qual := dax.NewTableQualifier(orgID, dbID) + qtid := dax.NewQualifiedTableID(qual, tableID) + + invalidQtid := dax.NewQualifiedTableID(qual, dax.TableID("0000000000000000")) + + ctx := context.Background() + + // Ensure that when using a Schemar not initiated with NewSchemar, the error + // handling works as expected. + t.Run("EmptyVersionStore", func(t *testing.T) { + s := inmem.VersionStore{} + + t.Run("GetShardsInvalid", func(t *testing.T) { + sh, ok, err := s.Shards(ctx, invalidQtid) + assert.NoError(t, err) + assert.False(t, ok) + assert.Nil(t, sh) + }) + + // Add new table. + assert.NoError(t, s.AddTable(ctx, qtid)) + }) + + t.Run("NewVersionStore", func(t *testing.T) { + s := inmem.NewVersionStore() + + // Add new table. + assert.NoError(t, s.AddTable(ctx, qtid)) + + t.Run("AddTableAgain", func(t *testing.T) { + err := s.AddTable(ctx, qtid) + if assert.Error(t, err) { + assert.True(t, errors.Is(err, dax.ErrTableIDExists)) + } + }) + + t.Run("AddShards", func(t *testing.T) { + err := s.AddShards(ctx, invalidQtid, dax.NewShard(1, 0)) + if assert.Error(t, err) { + assert.True(t, errors.Is(err, dax.ErrTableIDDoesNotExist)) + } + + { + _, ok, err := s.Shards(ctx, invalidQtid) + assert.NoError(t, err) + assert.False(t, ok) + } + + // Shards is empty if no shards have been added. + { + sh, ok, err := s.Shards(ctx, qtid) + assert.NoError(t, err) + assert.True(t, ok) + assert.Equal(t, sh, dax.Shards{}) + } + + // Add the first set of shards (with a duplicate (8)). + { + err := s.AddShards(ctx, qtid, + dax.NewShard(8, 0), + dax.NewShard(9, 0), + dax.NewShard(8, 0), + dax.NewShard(10, 0), + ) + assert.NoError(t, err) + } + + { + sh, ok, err := s.Shards(ctx, qtid) + assert.NoError(t, err) + assert.True(t, ok) + assert.Equal(t, dax.Shards{ + dax.NewShard(8, 0), + dax.NewShard(9, 0), + dax.NewShard(10, 0), + }, sh) + } + + // Add another set of shards (with one duplicate (11) and one + // existing (10)). + { + err := s.AddShards(ctx, qtid, + dax.NewShard(10, 0), + dax.NewShard(11, 0), + dax.NewShard(12, 0), + dax.NewShard(11, 0), + ) + assert.NoError(t, err) + } + + { + sh, ok, err := s.Shards(ctx, qtid) + assert.NoError(t, err) + assert.True(t, ok) + assert.Equal(t, dax.Shards{ + dax.NewShard(8, 0), + dax.NewShard(9, 0), + dax.NewShard(10, 0), + dax.NewShard(11, 0), + dax.NewShard(12, 0), + }, sh) + } + }) + + t.Run("RemoveTable", func(t *testing.T) { + shards, partitions, err := s.RemoveTable(ctx, qtid) + assert.NoError(t, err) + assert.Equal(t, dax.Partitions{}, partitions) + assert.Equal(t, dax.Shards{ + dax.NewShard(8, 0), + dax.NewShard(9, 0), + dax.NewShard(10, 0), + dax.NewShard(11, 0), + dax.NewShard(12, 0), + }, shards) + + // Make sure the table was removed. + shards, ok, err := s.Shards(ctx, qtid) + assert.NoError(t, err) + assert.False(t, ok) + assert.Nil(t, shards) + }) + }) + + t.Run("ErrorConditions", func(t *testing.T) { + t.Run("JustSchemar", func(t *testing.T) { + s := inmem.VersionStore{} + + shards, partitions, err := s.RemoveTable(ctx, qtid) + assert.Nil(t, shards) + assert.Nil(t, partitions) + if assert.Error(t, err) { + assert.True(t, errors.Is(err, dax.ErrTableIDDoesNotExist)) + } + }) + + t.Run("NewSchemar", func(t *testing.T) { + s := inmem.NewVersionStore() + + shards, partitions, err := s.RemoveTable(ctx, qtid) + assert.Nil(t, shards) + assert.Nil(t, partitions) + if assert.Error(t, err) { + assert.True(t, errors.Is(err, dax.ErrTableIDDoesNotExist)) + } + }) + }) +} diff --git a/dax/mds/api/openapi.yaml b/dax/mds/api/openapi.yaml new file mode 100644 index 000000000..bf4f84488 --- /dev/null +++ b/dax/mds/api/openapi.yaml @@ -0,0 +1,548 @@ +openapi: 3.0.3 + +info: + title: MDS + description: Metadata Services. + version: 0.0.0 + +paths: + /mds/health: + get: + summary: Health check endpoint. + description: Provides an endpoint to check the overall health of the MDS service. + operationId: GetHealth + responses: + 200: + description: Service is healthy. + + /mds/create-table: + post: + summary: Create a table. + description: Create a table based on the provided schema. + operationId: PostCreateTable + requestBody: + content: + application/json: + examples: + table: + $ref: '#/components/examples/Table' + schema: + $ref: '#/components/schemas/Table' + responses: + 200: + $ref: '#/components/responses/CreateTableResponse' + + /mds/drop-table: + post: + summary: Drop a table. + description: Drop a table based on the provided table name. + operationId: PostDropTable + requestBody: + content: + application/json: + example: + name: tbl + schema: + type: object + properties: + name: + type: string + responses: + 200: + description: Table was dropped. + + /mds/create-field: + post: + summary: Create a field. + description: Create a field based on the provided table and schema. + operationId: PostCreateField + requestBody: + content: + application/json: + example: + table: tbl + field: + name: a_string + type: string + options: + cacheType: ranked + cacheSize: 50000 + schema: + $ref: '#/components/schemas/TableField' + responses: + 200: + description: Field was created. + + /mds/drop-field: + post: + summary: Drop a field. + description: Drop a field based on the provided table and field name. + operationId: PostDropField + requestBody: + content: + application/json: + example: + table: tbl + field: fld + schema: + type: object + properties: + table: + type: string + field: + type: string + responses: + 200: + description: Field was dropped. + + /mds/table: + post: + summary: Get a table. + description: Get a table based on the provided table name. + operationId: PostTable + requestBody: + content: + application/json: + example: + name: tbl + schema: + type: object + properties: + name: + type: string + responses: + 200: + $ref: '#/components/responses/Table' + + /mds/tables: + post: + summary: Get a list of table. + description: Get a list of tables. If a filter is provided, only those tables will be included in the result. + operationId: PostTables + requestBody: + content: + application/json: + example: + names: + - tbl1 + - tbl2 + schema: + type: object + properties: + names: + type: array + items: + type: string + responses: + 200: + $ref: '#/components/responses/Tables' + + /mds/ingest-partition: + post: + summary: Request to ingest partition data. + description: Request to ingest (write) partition data. The address of the compute node responsible is returned. + operationId: PostIngestPartition + requestBody: + content: + application/json: + example: + table: tbl + partition: 7 + schema: + type: object + properties: + table: + type: string + partition: + type: integer + format: int32 + responses: + 200: + $ref: '#/components/responses/Address' + + /mds/ingest-shard: + post: + summary: Request to ingest shard data. + description: Request to ingest (write) shard data. The address of the compute node responsible is returned. + operationId: PostIngestShard + requestBody: + content: + application/json: + example: + table: tbl + shard: 12 + schema: + type: object + properties: + table: + type: string + shard: + type: integer + format: int64 + responses: + 200: + $ref: '#/components/responses/Address' + + /mds/snapshot/shard-data: + post: + summary: Request to snapshot shard data. + description: Request to snapshot shard data. + operationId: PostSnapshotShardData + requestBody: + content: + application/json: + example: + table: tbl + shard: 12 + schema: + type: object + properties: + table: + type: string + shard: + type: integer + format: int64 + responses: + 200: + description: Shard snapshot was successful. + + /mds/snapshot/table-keys: + post: + summary: Request to snapshot table keys. + description: Request to snapshot table keys. + operationId: PostSnapshotTableKeys + requestBody: + content: + application/json: + example: + table: tbl + partition: 7 + schema: + type: object + properties: + table: + type: string + partition: + type: integer + format: int32 + responses: + 200: + description: Table keys snapshot was successful. + + /mds/snapshot/field-keys: + post: + summary: Request to snapshot field keys. + description: Request to snapshot field keys. + operationId: PostSnapshotFieldKeys + requestBody: + content: + application/json: + example: + table: tbl + field: fld + schema: + type: object + properties: + table: + type: string + field: + type: string + responses: + 200: + description: Field keys snapshot was successful. + + /mds/register-node: + post: + summary: Register node. + description: Register a node with MDS. + operationId: PostRegisterNode + requestBody: + content: + application/json: + example: + address: 10.0.0.1:8000 + roleTypes: + - compute + - translate + schema: + type: object + properties: + address: + type: string + roleTypes: + type: array + items: + types: string + responses: + 200: + description: Node registration was successful. + + /mds/deregister-nodes: + post: + summary: Deregister nodes. + description: Deregister nodes with MDS. + operationId: PostDeregisterNodes + requestBody: + content: + application/json: + example: + address: 10.0.0.1:8000 + schema: + type: object + properties: + address: + type: string + responses: + 200: + description: Node deregistration was successful. + + /mds/compute-nodes: + post: + summary: Get compute nodes. + description: Get the compute nodes responsible for the given shards. + operationId: PostComputeNodes + requestBody: + content: + application/json: + example: + table: tbl + shards: + - 10 + - 11 + - 12 + isWrite: false + schema: + type: object + properties: + table: + type: string + shards: + type: array + items: + type: integer + format: int64 + isWrite: + type: boolean + responses: + 200: + $ref: '#/components/responses/ComputeNodes' + + /mds/translate-nodes: + post: + summary: Get translate nodes. + description: Get the translate nodes responsible for the given partitions. + operationId: PostTranslateNodes + requestBody: + content: + application/json: + example: + table: tbl + partitions: + - 6 + - 7 + isWrite: false + schema: + type: object + properties: + table: + type: string + partitions: + type: array + items: + type: integer + format: int32 + isWrite: + type: boolean + responses: + 200: + $ref: '#/components/responses/TranslateNodes' + +components: + responses: + CreateTableResponse: + description: Placeholder response. + content: + application/json: + schema: + type: object + + Address: + description: Single node address. + content: + application/json: + schema: + type: object + properties: + address: + type: string + + Table: + description: Table response. + content: + application/json: + schema: + $ref: '#/components/schemas/Table' + + Tables: + description: Tables response. + content: + application/json: + schema: + type: array + items: + $ref: '#/components/schemas/Table' + + ComputeNodes: + description: Compute nodes response. + content: + application/json: + schema: + type: array + items: + $ref: '#/components/schemas/ComputeNode' + + TranslateNodes: + description: Translate nodes response. + content: + application/json: + schema: + type: array + items: + $ref: '#/components/schemas/TranslateNode' + + schemas: + ComputeNode: + type: object + properties: + address: + type: string + table: + type: string + shards: + type: array + items: + type: integer + format: int64 + + TranslateNode: + type: object + properties: + address: + type: string + table: + type: string + partitions: + type: array + items: + type: integer + format: int32 + + Table: + type: object + properties: + name: + type: string + fields: + type: array + items: + $ref: '#/components/schemas/Field' + partitionN: + type: integer + format: int32 + + TableField: + type: object + properties: + table: + type: string + field: + $ref: '#/components/schemas/Field' + + Field: + type: object + properties: + name: + type: string + type: + type: string + enum: + - bool + - decimal + - id + - idset + - int + - string + - stringset + - timestamp + options: + type: object + properties: + min: + type: integer + format: int64 + max: + type: integer + format: int64 + scale: + type: integer + format: int64 + minimum: 0 + noStandardView: + type: boolean + cacheType: + type: string + cacheSize: + type: integer + format: int32 + timeUnit: + type: string + epoch: + type: string + format: date-time + timeQuantum: + type: string + ttl: + type: string + foreignIndex: + type: string + + examples: + Table: + name: tbl + fields: + - name: _id + type: string + - name: a_bool + type: bool + - name: an_id + type: id + options: + cacheType: ranked + cacheSize: 50000 + - name: an_id_set + type: idset + options: + cacheType: ranked + cacheSize: 50000 + - name: a_string + type: string + options: + cacheType: ranked + cacheSize: 50000 + - name: a_string_set + type: stringset + options: + cacheType: ranked + cacheSize: 50000 + - name: an_int + type: int + options: + min: -100 + max: 500 + - name: a_decimal + type: decimal + options: + min: -10.24 + max: 50.75 + scale: 2 + partitionN: 16 diff --git a/dax/mds/client/client.go b/dax/mds/client/client.go new file mode 100644 index 000000000..c99affd89 --- /dev/null +++ b/dax/mds/client/client.go @@ -0,0 +1,478 @@ +// Package client is an HTTP client for MDS. +package client + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "log" + "net/http" + + fb "github.com/molecula/featurebase/v3" + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/mds/controller" + mdshttp "github.com/molecula/featurebase/v3/dax/mds/http" + "github.com/molecula/featurebase/v3/errors" +) + +const ( + defaultScheme = "http" + defaultPath = "/mds" +) + +// Ensure type implements interface. +var _ fb.MDS = (*Client)(nil) + +// Client is an HTTP client that operates on the MDS endpoints exposed by the +// main MDS service. +type Client struct { + address dax.Address +} + +// New returns a new instance of Client. +func New(address dax.Address) *Client { + return &Client{ + address: address, + } +} + +// Health returns true if the client address returns status OK at its /health +// endpoint. +func (c *Client) Health() bool { + url := fmt.Sprintf("%s%s/health", c.address.WithScheme(defaultScheme), defaultPath) + + if resp, err := http.Get(url); err != nil { + return false + } else if resp.StatusCode != http.StatusOK { + return false + } + + return true +} + +func (c *Client) Table(ctx context.Context, qtid dax.QualifiedTableID) (*dax.QualifiedTable, error) { + url := fmt.Sprintf("%s%s/table", c.address.WithScheme(defaultScheme), defaultPath) + + // Encode the request. + postBody, err := json.Marshal(qtid) + if err != nil { + return nil, errors.Wrap(err, "marshalling post request") + } + responseBody := bytes.NewBuffer(postBody) + + // Post the request. + log.Printf("POST table request: url: %s", url) + resp, err := http.Post(url, "application/json", responseBody) + if err != nil { + return nil, errors.Wrap(err, "posting table request") + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(resp.Body) + return nil, errors.Errorf("status code: %d: %s", resp.StatusCode, b) + } + + var qtable *dax.QualifiedTable + if err := json.NewDecoder(resp.Body).Decode(&qtable); err != nil { + return nil, errors.Wrap(err, "reading response body") + } + + return qtable, nil +} + +func (c *Client) TableID(ctx context.Context, qual dax.TableQualifier, name dax.TableName) (dax.QualifiedTableID, error) { + url := fmt.Sprintf("%s%s/table-id", c.address.WithScheme(defaultScheme), defaultPath) + + dflt := dax.QualifiedTableID{} + + req := dax.QualifiedTableID{ + TableQualifier: qual, + Name: name, + } + + // Encode the request. + postBody, err := json.Marshal(req) + if err != nil { + return dflt, errors.Wrap(err, "marshalling post request") + } + requestBody := bytes.NewBuffer(postBody) + + // Post the request. + resp, err := http.Post(url, "application/json", requestBody) + if err != nil { + return dflt, errors.Wrap(err, "posting table-id request") + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(resp.Body) + return dflt, errors.Errorf("status code: %d: %s", resp.StatusCode, b) + } + + var qtid dax.QualifiedTableID + if err := json.NewDecoder(resp.Body).Decode(&qtid); err != nil { + return dflt, errors.Wrap(err, "reading response body") + } + + return qtid, nil +} + +func (c *Client) Tables(ctx context.Context, qual dax.TableQualifier, ids ...dax.TableID) ([]*dax.QualifiedTable, error) { + url := fmt.Sprintf("%s%s/tables", c.address.WithScheme(defaultScheme), defaultPath) + + req := mdshttp.TablesRequest{ + OrganizationID: qual.OrganizationID, + DatabaseID: qual.DatabaseID, + TableIDs: ids, + } + + // Encode the request. + postBody, err := json.Marshal(req) + if err != nil { + return nil, errors.Wrap(err, "marshalling post request") + } + responseBody := bytes.NewBuffer(postBody) + + // Post the request. + resp, err := http.Post(url, "application/json", responseBody) + if err != nil { + return nil, errors.Wrap(err, "posting tables request") + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(resp.Body) + return nil, errors.Errorf("status code: %d: %s", resp.StatusCode, b) + } + + var qtables []*dax.QualifiedTable + if err := json.NewDecoder(resp.Body).Decode(&qtables); err != nil { + return nil, errors.Wrap(err, "reading response body") + } + + return qtables, nil +} + +func (c *Client) CreateTable(ctx context.Context, qtbl *dax.QualifiedTable) error { + url := fmt.Sprintf("%s%s/create-table", c.address.WithScheme(defaultScheme), defaultPath) + + // Encode the request. + postBody, err := json.Marshal(qtbl) + if err != nil { + return errors.Wrap(err, "marshalling post request") + } + responseBody := bytes.NewBuffer(postBody) + + // Post the request. + resp, err := http.Post(url, "application/json", responseBody) + if err != nil { + return errors.Wrap(err, "posting create table request") + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(resp.Body) + return errors.Errorf("status code: %d: %s", resp.StatusCode, b) + } + + return nil +} + +func (c *Client) DropTable(ctx context.Context, qtid dax.QualifiedTableID) error { + url := fmt.Sprintf("%s%s/drop-table", c.address.WithScheme(defaultScheme), defaultPath) + + // Encode the request. + postBody, err := json.Marshal(qtid) + if err != nil { + return errors.Wrap(err, "marshalling post request") + } + responseBody := bytes.NewBuffer(postBody) + + // Post the request. + resp, err := http.Post(url, "application/json", responseBody) + if err != nil { + return errors.Wrap(err, "posting drop table request") + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(resp.Body) + return errors.Errorf("status code: %d: %s", resp.StatusCode, b) + } + + return nil +} + +func (c *Client) CreateField(ctx context.Context, qtid dax.QualifiedTableID, fld *dax.Field) error { + url := fmt.Sprintf("%s%s/create-field", c.address.WithScheme(defaultScheme), defaultPath) + + req := mdshttp.CreateFieldRequest{ + TableKey: qtid.Key(), + Field: fld, + } + // Encode the request. + postBody, err := json.Marshal(req) + if err != nil { + return errors.Wrap(err, "marshalling post request") + } + responseBody := bytes.NewBuffer(postBody) + + // Post the request. + resp, err := http.Post(url, "application/json", responseBody) + if err != nil { + return errors.Wrap(err, "posting create field request") + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(resp.Body) + return errors.Errorf("status code: %d: %s", resp.StatusCode, b) + } + + return nil +} + +func (c *Client) DropField(ctx context.Context, qtid dax.QualifiedTableID, fldName dax.FieldName) error { + url := fmt.Sprintf("%s%s/drop-field", c.address.WithScheme(defaultScheme), defaultPath) + + // Encode the request. + req := mdshttp.DropFieldRequest{ + Table: qtid, + Field: fldName, + } + + postBody, err := json.Marshal(req) + if err != nil { + return errors.Wrap(err, "marshalling post request") + } + responseBody := bytes.NewBuffer(postBody) + + // Post the request. + resp, err := http.Post(url, "application/json", responseBody) + if err != nil { + return errors.Wrap(err, "posting drop field request") + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(resp.Body) + return errors.Errorf("status code: %d: %s", resp.StatusCode, b) + } + + return nil +} + +func (c *Client) IngestShard(ctx context.Context, qtid dax.QualifiedTableID, shard dax.ShardNum) (dax.Address, error) { + url := fmt.Sprintf("%s%s/ingest-shard", c.address.WithScheme(defaultScheme), defaultPath) + + var host dax.Address + + req := &mdshttp.IngestShardRequest{ + Table: qtid, + Shard: shard, + } + + // Encode the request. + postBody, err := json.Marshal(req) + if err != nil { + return host, errors.Wrap(err, "marshalling post request") + } + responseBody := bytes.NewBuffer(postBody) + + // Post the request. + resp, err := http.Post(url, "application/json", responseBody) + if err != nil { + return host, errors.Wrap(err, "posting ingest-shard request") + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(resp.Body) + return host, errors.Errorf("status code: %d: %s", resp.StatusCode, b) + } + + var isr *mdshttp.IngestShardResponse + if err := json.NewDecoder(resp.Body).Decode(&isr); err != nil { + return host, errors.Wrap(err, "reading response body") + } + + return isr.Address, nil +} + +func (c *Client) IngestPartition(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum) (dax.Address, error) { + url := fmt.Sprintf("%s%s/ingest-partition", c.address.WithScheme(defaultScheme), defaultPath) + + var host dax.Address + + req := &mdshttp.IngestPartitionRequest{ + Table: qtid, + Partition: partition, + } + + // Encode the request. + postBody, err := json.Marshal(req) + if err != nil { + return host, errors.Wrap(err, "marshalling post request") + } + responseBody := bytes.NewBuffer(postBody) + + // Post the request. + resp, err := http.Post(url, "application/json", responseBody) + if err != nil { + return host, errors.Wrap(err, "posting ingest-partition request") + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(resp.Body) + return host, errors.Errorf("status code: %d: %s", resp.StatusCode, b) + } + + var isr *mdshttp.IngestPartitionResponse + if err := json.NewDecoder(resp.Body).Decode(&isr); err != nil { + return host, errors.Wrap(err, "reading response body") + } + + return isr.Address, nil +} + +func (c *Client) ComputeNodes(ctx context.Context, qtid dax.QualifiedTableID, shards ...dax.ShardNum) ([]controller.ComputeNode, error) { + url := fmt.Sprintf("%s%s/compute-nodes", c.address.WithScheme(defaultScheme), defaultPath) + log.Printf("ComputeNodes url: %s", url) + + var nodes []controller.ComputeNode + + req := &mdshttp.ComputeNodesRequest{ + Table: qtid, + Shards: shards, + } + + // Encode the request. + postBody, err := json.Marshal(req) + if err != nil { + return nodes, errors.Wrap(err, "marshalling post request") + } + responseBody := bytes.NewBuffer(postBody) + + // Post the request. + resp, err := http.Post(url, "application/json", responseBody) + if err != nil { + return nodes, errors.Wrap(err, "posting compute-nodes request") + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(resp.Body) + return nodes, errors.Errorf("status code: %d: %s", resp.StatusCode, b) + } + + var cnr *mdshttp.ComputeNodesResponse + if err := json.NewDecoder(resp.Body).Decode(&cnr); err != nil { + return nodes, errors.Wrap(err, "reading response body") + } + + return cnr.ComputeNodes, nil +} + +func (c *Client) TranslateNodes(ctx context.Context, qtid dax.QualifiedTableID, partitions ...dax.PartitionNum) ([]controller.TranslateNode, error) { + url := fmt.Sprintf("%s%s/translate-nodes", c.address.WithScheme(defaultScheme), defaultPath) + log.Printf("TranslateNodes url: %s", url) + + var nodes []controller.TranslateNode + + req := &mdshttp.TranslateNodesRequest{ + Table: qtid, + Partitions: partitions, + } + + // Encode the request. + postBody, err := json.Marshal(req) + if err != nil { + return nodes, errors.Wrap(err, "marshalling post request") + } + responseBody := bytes.NewBuffer(postBody) + + // Post the request. + resp, err := http.Post(url, "application/json", responseBody) + if err != nil { + return nodes, errors.Wrap(err, "posting translate-nodes request") + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(resp.Body) + return nodes, errors.Errorf("status code: %d: %s", resp.StatusCode, b) + } + + var cnr *mdshttp.TranslateNodesResponse + if err := json.NewDecoder(resp.Body).Decode(&cnr); err != nil { + return nodes, errors.Wrap(err, "reading response body") + } + + return cnr.TranslateNodes, nil +} + +func (c *Client) RegisterNode(ctx context.Context, node *dax.Node) error { + url := fmt.Sprintf("%s%s/register-node", c.address.WithScheme(defaultScheme), defaultPath) + log.Printf("RegisterNode url: %s", url) + + req := &mdshttp.RegisterNodeRequest{ + Address: node.Address, + RoleTypes: node.RoleTypes, + } + + // Encode the request. + postBody, err := json.Marshal(req) + if err != nil { + return errors.Wrap(err, "marshalling post request") + } + responseBody := bytes.NewBuffer(postBody) + + // Post the request. + resp, err := http.Post(url, "application/json", responseBody) + if err != nil { + return errors.Wrap(err, "posting translate-nodes request") + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(resp.Body) + return errors.Errorf("status code: %d: %s", resp.StatusCode, b) + } + + return nil +} + +func (c *Client) CheckInNode(ctx context.Context, node *dax.Node) error { + url := fmt.Sprintf("%s%s/check-in-node", c.address.WithScheme(defaultScheme), defaultPath) + log.Printf("CheckInNode url: %s", url) + + req := &mdshttp.CheckInNodeRequest{ + Address: node.Address, + RoleTypes: node.RoleTypes, + } + + // Encode the request. + postBody, err := json.Marshal(req) + if err != nil { + return errors.Wrap(err, "marshalling post request") + } + responseBody := bytes.NewBuffer(postBody) + + // Post the request. + resp, err := http.Post(url, "application/json", responseBody) + if err != nil { + return errors.Wrap(err, "posting translate-nodes request") + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(resp.Body) + return errors.Errorf("status code: %d: %s", resp.StatusCode, b) + } + + return nil +} diff --git a/dax/mds/controller/alpha/director.go b/dax/mds/controller/alpha/director.go new file mode 100644 index 000000000..ef8395d63 --- /dev/null +++ b/dax/mds/controller/alpha/director.go @@ -0,0 +1,125 @@ +// Package alpha contains inter-service implemenations of interfaces. +package alpha + +import ( + "context" + "encoding/json" + + featurebase "github.com/molecula/featurebase/v3" + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/mds/controller" + "github.com/molecula/featurebase/v3/errors" + featurebaseserver "github.com/molecula/featurebase/v3/server" +) + +// Ensure type implements interface. +var _ controller.Director = (*Director)(nil) + +// Director is a direct, service-to-service implementation of the Director +// interface. +type Director struct { + computers map[dax.Address]*featurebaseserver.Command +} + +func NewDirector() *Director { + return &Director{ + computers: make(map[dax.Address]*featurebaseserver.Command), + } +} + +func (d *Director) AddCmd(addr dax.Address, cmd *featurebaseserver.Command) error { + if cmd == nil { + return errors.New(errors.ErrUncoded, "cannot add nil cmd to director") + } + d.computers[addr] = cmd + return nil +} + +func (d *Director) api(addr dax.Address) (*featurebase.API, error) { + cmd, found := d.computers[addr] + if !found { + // Address not registered with the Director. + return nil, errors.New(errors.ErrUncoded, "cmd not registered with director") + } + + api := cmd.API + if api == nil { + // Command does not have an API. + return nil, errors.New(errors.ErrUncoded, "cmd does not have an api") + } + + return api, nil +} + +func (d *Director) SendDirective(ctx context.Context, dir *dax.Directive) error { + api, err := d.api(dir.Address) + if err != nil { + return errors.Wrap(err, "getting api from director") + } + + ndir, err := marshalUnmarshal(dir) + if err != nil { + return errors.Wrap(err, "marshalUnmarshal") + } + + return api.Directive(ctx, ndir) +} + +func (d *Director) SendSnapshotShardDataRequest(ctx context.Context, req *dax.SnapshotShardDataRequest) error { + api, err := d.api(req.Address) + if err != nil { + return errors.Wrap(err, "getting api from director") + } + + nreq, err := marshalUnmarshal(req) + if err != nil { + return errors.Wrap(err, "marshalUnmarshal") + } + + return api.SnapshotShardData(ctx, nreq) +} + +func (d *Director) SendSnapshotTableKeysRequest(ctx context.Context, req *dax.SnapshotTableKeysRequest) error { + api, err := d.api(req.Address) + if err != nil { + return errors.Wrap(err, "getting api from director") + } + + nreq, err := marshalUnmarshal(req) + if err != nil { + return errors.Wrap(err, "marshalUnmarshal") + } + + return api.SnapshotTableKeys(ctx, nreq) +} + +func (d *Director) SendSnapshotFieldKeysRequest(ctx context.Context, req *dax.SnapshotFieldKeysRequest) error { + api, err := d.api(req.Address) + if err != nil { + return errors.Wrap(err, "getting api from director") + } + + nreq, err := marshalUnmarshal(req) + if err != nil { + return errors.Wrap(err, "marshalUnmarshal") + } + + return api.SnapshotFieldKeys(ctx, nreq) +} + +// marshalUnmarshal simply marshals anything to json, and then +// unmarshals it. This might seem a bit silly. The reason it exists is +// to exercise the same encode/decode logic that we'd need to if we +// were traversing the network, and guarantee that we aren't sharing +// pointers across API boundaries. +func marshalUnmarshal[K any](a K) (K, error) { + var newA K + abytes, err := json.Marshal(a) + if err != nil { + return newA, errors.Wrap(err, "marshaling directive") + } + if err := json.Unmarshal(abytes, &newA); err != nil { + return newA, errors.Wrap(err, "unmarshaling directive") + } + return newA, nil +} diff --git a/dax/mds/controller/balancer.go b/dax/mds/controller/balancer.go new file mode 100644 index 000000000..5c2457ab8 --- /dev/null +++ b/dax/mds/controller/balancer.go @@ -0,0 +1,69 @@ +package controller + +import ( + "context" + "fmt" + + "github.com/molecula/featurebase/v3/dax" +) + +type Balancer interface { + AddWorker(ctx context.Context, worker fmt.Stringer) ([]dax.WorkerDiff, error) + RemoveWorker(ctx context.Context, worker fmt.Stringer) ([]dax.WorkerDiff, error) + AddJob(ctx context.Context, job fmt.Stringer) ([]dax.WorkerDiff, error) + RemoveJob(ctx context.Context, job fmt.Stringer) ([]dax.WorkerDiff, error) + Balance(ctx context.Context) ([]dax.WorkerDiff, error) + CurrentState(ctx context.Context) ([]dax.WorkerInfo, error) + WorkerState(ctx context.Context, worker dax.Worker) (dax.WorkerInfo, error) + WorkersForJobs(ctx context.Context, jobs []dax.Job) ([]dax.WorkerInfo, error) + + // WorkersForJobPrefix returns all workers and their job + // assignments which start with `prefix` for all jobs that start + // with `prefix`. If there are free jobs that start with `prefix` + // an error is returned. + // + // The motivating use case is getting all workers for a particular + // table so we can execute a query that will hit every shard in a + // table. If there are jobs representing shards in that table + // which are not assigned to any worker, that means the query + // would return incomplete data, so we want to error. + WorkersForJobPrefix(ctx context.Context, prefix string) ([]dax.WorkerInfo, error) +} + +// Ensure type implements interface. +var _ Balancer = (*NopBalancer)(nil) + +// NopBalancer is a no-op implementation of the Balancer interface. +type NopBalancer struct{} + +func NewNopBalancer() *NopBalancer { + return &NopBalancer{} +} + +func (b *NopBalancer) AddWorker(ctx context.Context, worker fmt.Stringer) ([]dax.WorkerDiff, error) { + return []dax.WorkerDiff{}, nil +} +func (b *NopBalancer) RemoveWorker(ctx context.Context, worker fmt.Stringer) ([]dax.WorkerDiff, error) { + return []dax.WorkerDiff{}, nil +} +func (b *NopBalancer) AddJob(ctx context.Context, job fmt.Stringer) ([]dax.WorkerDiff, error) { + return []dax.WorkerDiff{}, nil +} +func (b *NopBalancer) RemoveJob(ctx context.Context, job fmt.Stringer) ([]dax.WorkerDiff, error) { + return []dax.WorkerDiff{}, nil +} +func (b *NopBalancer) Balance(ctx context.Context) ([]dax.WorkerDiff, error) { + return []dax.WorkerDiff{}, nil +} +func (b *NopBalancer) CurrentState(ctx context.Context) ([]dax.WorkerInfo, error) { + return []dax.WorkerInfo{}, nil +} +func (b *NopBalancer) WorkerState(ctx context.Context, worker dax.Worker) (dax.WorkerInfo, error) { + return dax.WorkerInfo{}, nil +} +func (b *NopBalancer) WorkersForJobs(ctx context.Context, jobs []dax.Job) ([]dax.WorkerInfo, error) { + return []dax.WorkerInfo{}, nil +} +func (b *NopBalancer) WorkersForJobPrefix(ctx context.Context, prefix string) ([]dax.WorkerInfo, error) { + return []dax.WorkerInfo{}, nil +} diff --git a/dax/mds/controller/config.go b/dax/mds/controller/config.go new file mode 100644 index 000000000..0c3ac19b1 --- /dev/null +++ b/dax/mds/controller/config.go @@ -0,0 +1,29 @@ +package controller + +import ( + "time" + + "github.com/molecula/featurebase/v3/dax/boltdb" + "github.com/molecula/featurebase/v3/dax/mds/schemar" + "github.com/molecula/featurebase/v3/logger" +) + +type NewBalancerFn func(string, logger.Logger) Balancer + +type Config struct { + Director Director + Schemar schemar.Schemar + ComputeBalancer Balancer + TranslateBalancer Balancer + + StorageMethod string + BoltDB *boltdb.DB + + // RegistrationBatchTimeout is the time that the controller will + // wait after a node registers itself to see if any more nodes + // will register before sending out directives to all nodes which + // have been registered. + RegistrationBatchTimeout time.Duration + + Logger logger.Logger +} diff --git a/dax/mds/controller/controller.go b/dax/mds/controller/controller.go new file mode 100644 index 000000000..c6f130fc3 --- /dev/null +++ b/dax/mds/controller/controller.go @@ -0,0 +1,1892 @@ +// Package controller provides the core Controller struct. +package controller + +import ( + "context" + "fmt" + "sort" + "sync" + "time" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/boltdb" + "github.com/molecula/featurebase/v3/dax/mds/schemar" + "github.com/molecula/featurebase/v3/errors" + "github.com/molecula/featurebase/v3/logger" + "golang.org/x/sync/errgroup" +) + +type Controller struct { + // mu is primarily to protect against conflicting reads/writes to the nodes + // map. The balancers map is currently never written to after + // initialization. + mu sync.RWMutex + + // versionStore + versionStore dax.VersionStore + + // Schemar used by the controller to get table information. The controller + // should NOT call Schemar methods which modify data. Schema mutations are + // made outside of the controller (at this point that happens in MDS). + Schemar schemar.Schemar + + // nodes is the map of nodes, by address, which have registered with the + // controller. + nodeService dax.NodeService + + ComputeBalancer Balancer + TranslateBalancer Balancer + + // Director is used to send directives to computer workers. + Director Director + + // poller is used to notify a Poller if nodes have been added or removed. + poller dax.AddressManager + + directiveVersion dax.DirectiveVersion + + registrationBatchTimeout time.Duration + nodeChan chan *dax.Node + stopping chan struct{} + + logger logger.Logger +} + +var supportedRoleTypes []dax.RoleType = []dax.RoleType{ + dax.RoleTypeCompute, + dax.RoleTypeTranslate, +} + +// New returns a new instance of Controller with default values. +func New(cfg Config) *Controller { + c := &Controller{ + Schemar: schemar.NewNopSchemar(), + + ComputeBalancer: cfg.ComputeBalancer, + TranslateBalancer: cfg.TranslateBalancer, + + Director: NewNopDirector(), + + poller: dax.NewNopAddressManager(), + + logger: logger.NopLogger, + + nodeChan: make(chan *dax.Node, 10), + } + + if cfg.Logger != nil { + c.logger = cfg.Logger + } + + switch cfg.StorageMethod { + case "boltdb": + if err := cfg.BoltDB.InitializeBuckets(boltdb.VersionStoreBuckets...); err != nil { + c.logger.Panicf("initializing version store buckets: %v", err) + } + c.versionStore = boltdb.NewVersionStore(cfg.BoltDB, c.logger) + + if err := cfg.BoltDB.InitializeBuckets(boltdb.NodeServiceBuckets...); err != nil { + c.logger.Panicf("initializing node service buckets: %v", err) + } + c.nodeService = boltdb.NewNodeService(cfg.BoltDB, c.logger) + + if err := cfg.BoltDB.InitializeBuckets(boltdb.DirectiveBuckets...); err != nil { + c.logger.Panicf("initializing directive buckets: %v", err) + } + c.directiveVersion = boltdb.NewDirectiveVersion(cfg.BoltDB) + default: + c.logger.Panicf("storage method '%s' unsupported. (hint: try boltdb)", cfg.StorageMethod) + } + + if cfg.Director != nil { + c.Director = cfg.Director + } + if cfg.Schemar != nil { + c.Schemar = cfg.Schemar + } + + c.registrationBatchTimeout = cfg.RegistrationBatchTimeout + + return c +} + +// Run starts the node registration goroutine. +func (c *Controller) Run() error { + go c.nodeRegistrationRoutine(c.nodeChan, c.registrationBatchTimeout) + + return nil +} + +// Stop stops the node registration routine. +func (c *Controller) Stop() { + close(c.stopping) +} + +func (c *Controller) balancerForRole(rt dax.RoleType) (Balancer, error) { + var bal Balancer + if rt == dax.RoleTypeCompute { + bal = c.ComputeBalancer + } else if rt == dax.RoleTypeTranslate { + bal = c.TranslateBalancer + } else { + return nil, errors.Errorf("unknown role type: '%s'", rt) + } + return bal, nil +} + +// RegisterNodes adds nodes to the controller's list of registered +// nodes. +func (c *Controller) RegisterNodes(ctx context.Context, nodes ...*dax.Node) error { + c.logger.Printf("c.RegisterNodes(): %+v", nodes) + c.mu.Lock() + defer c.mu.Unlock() + + // Validate input. + for _, n := range nodes { + if n.Address == "" { + return NewErrNodeKeyInvalid(n.Address) + } + if len(n.RoleTypes) == 0 { + return NewErrRoleTypeInvalid(dax.RoleType("")) + } + for _, v := range n.RoleTypes { + if !dax.RoleTypes(supportedRoleTypes).Contains(v) { + return NewErrRoleTypeInvalid(v) + } + } + } + + // workerSet maintains the set of workers which have a job assignment change + // and therefore need to be sent an updated Directive. + workerSet := NewAddressSet() + + // Create node if we don't already have it + for _, n := range nodes { + if node, _ := c.nodeService.ReadNode(ctx, n.Address); node == nil { + if err := c.nodeService.CreateNode(ctx, n.Address, n); err != nil { + return errors.Wrapf(err, "creating node: %s", n.Address) + } + + // Add the node to the workerSet so that it receives a directive. + // Even if there is currently no data for this worker (i.e. it + // doesn't result in a diffByAddr entry below), we still want to + // send it a "reset" directive so that in the off chance it has some + // local data, that data gets removed. + workerSet.Add(n.Address) + } + } + + // diffByAddr keeps track of the diffs that have been applied to each + // specific address. + diffByAddr := make(map[dax.Address]dax.WorkerDiff) + + for _, n := range nodes { + for _, rt := range n.RoleTypes { + balancer, err := c.balancerForRole(rt) + if err != nil { + return errors.Wrap(err, "getting balancer") + } + adiffs, err := balancer.AddWorker(ctx, n.Address) + if err != nil { + return errors.Wrap(err, "adding worker") + } + + // Rebalance so existing jobs can be spread evenly across all nodes, + // including the node being registered. + bdiffs, err := balancer.Balance(ctx) + if err != nil { + return errors.Wrap(err, "balancing") + } + for _, diff := range append(adiffs, bdiffs...) { + existingDiff, ok := diffByAddr[dax.Address(diff.WorkerID)] + if !ok { + existingDiff.WorkerID = diff.WorkerID + } + existingDiff.Add(diff) + diffByAddr[dax.Address(diff.WorkerID)] = existingDiff + } + } + } + + // Add any worker which has a diff to the workerSet so that it receives a + // directive. + for addr := range diffByAddr { + workerSet.Add(addr) + } + + addrs := []dax.Address{} + for _, n := range nodes { + addrs = append(addrs, n.Address) + } + + // Tell the poller about the new nodes. + if err := c.poller.AddAddresses(ctx, addrs...); err != nil { + return NewErrInternal(err.Error()) + } + + // No need to send directives if the workerSet is empty. + if len(workerSet) == 0 { + return nil + } + + // Convert the slice of addresses into a slice of addressMethod containing + // the appropriate method. + addressMethods := applyAddressMethod(workerSet.SortedSlice(), dax.DirectiveMethodDiff) + + // For the addresses which are being added, set their method to "reset". + for i := range addressMethods { + for j := range nodes { + if addressMethods[i].address == nodes[j].Address { + addressMethods[i].method = dax.DirectiveMethodReset + } + } + } + + // Get the current job assignments for this worker and send that to the node + // as a Directive. + if err := c.sendDirectives(ctx, addressMethods...); err != nil { + return NewErrDirectiveSendFailure(err.Error()) + } + + return nil +} + +// RegisterNode adds a node to the controller's list of registered +// nodes. It makes no guarantees about when the node will actually be +// used for anything or assigned any jobs. +func (c *Controller) RegisterNode(ctx context.Context, n *dax.Node) error { + c.mu.Lock() + defer c.mu.Unlock() + + // Validate input. + if n.Address == "" { + return NewErrNodeKeyInvalid(n.Address) + } + if len(n.RoleTypes) == 0 { + return NewErrRoleTypeInvalid(dax.RoleType("")) + } + for _, v := range n.RoleTypes { + if !dax.RoleTypes(supportedRoleTypes).Contains(v) { + return NewErrRoleTypeInvalid(v) + } + } + + if node, _ := c.nodeService.ReadNode(ctx, n.Address); node != nil { + return nil + } + + c.nodeChan <- n + + return nil +} + +// CheckInNode handles a "check-in" from a compute node. These come +// periodically, and if the controller already knows about the compute node, it +// can simply no-op. If, however, the controller is not aware of the node +// checking in, then that probably means that the poller has removed that node +// from its list (perhaps due to a network fault) and therefore the node needs +// to be re-registered. +func (c *Controller) CheckInNode(ctx context.Context, n *dax.Node) error { + c.mu.RLock() + defer c.mu.RUnlock() + + // If we already know about this node, just no-op. In the future, we may + // want this check-in payload to include things like the compute node's + // Directive; then we could check that the compute node is actually doing + // what we expect it to be doing. But for now, we're just checking that we + // know about the compute node at all. + if node, _ := c.nodeService.ReadNode(ctx, n.Address); node != nil { + return nil + } + + c.nodeChan <- n + + return nil +} + +// DeregisterNodes removes nodes from the controller's list of registered nodes. +// It sends directives to the removed nodes, but ignores errors. +func (c *Controller) DeregisterNodes(ctx context.Context, addresses ...dax.Address) error { + c.mu.Lock() + defer c.mu.Unlock() + + // workerSet maintains the set of workers which have a job assignment change + // and therefore need to be sent an updated Directive. + workerSet := NewAddressSet() + + // diffByAddr keeps track of the diffs that have been applied to each + // specific address. + diffByAddr := make(map[dax.Address]dax.WorkerDiff) + + for _, address := range addresses { + // Add the removed node to the workerSet so that it receives a + // directive. Even if there is currently no data for the worker (i.e. it + // doesn't result in a diffByAddr entry below), we still want to send it + // a "reset" directive so that in the off chance it has some local data, + // that data gets removed. + // TODO(tlt): see below where we actually REMOVE this. We need to + // address this confusion. + // workerSet.Add(address) + + // Ensure the host:port is currently registered. + n, err := c.nodeService.ReadNode(ctx, address) + if err != nil { + return errors.Wrapf(err, "reading the node for address: %s", address) + } + for _, rt := range n.RoleTypes { + balancer, err := c.balancerForRole(rt) + if err != nil { + c.logger.Printf("Unsupported role type in DeregisterNode: '%s'", rt) + // Skip any role types which aren't currently supported by a balancer. + continue + } + rdiffs, err := balancer.RemoveWorker(ctx, address) + if err != nil { + return errors.Wrap(err, "removing worker") + } + + // Rebalance so any jobs that were assigned to the node being deregistered + // get assigned to another node. + bdiffs, err := balancer.Balance(ctx) + if err != nil { + return errors.Wrap(err, "balancing") + } + // we assume that the job names are different between the + // different role types so we don't have to track each + // role separately which would be annoying. + for _, diff := range append(rdiffs, bdiffs...) { + existingDiff, ok := diffByAddr[dax.Address(diff.WorkerID)] + if !ok { + existingDiff.WorkerID = diff.WorkerID + } + existingDiff.Add(diff) + diffByAddr[dax.Address(diff.WorkerID)] = existingDiff + } + } + } + + for addr := range diffByAddr { + workerSet.Add(addr) + } + + // Don't send a Directive to the removed nodes after all. + // TODO(tlt): we have to do this because otherwise the send request hangs + // while holding a mu.Lock on Controller. + for _, addr := range addresses { + workerSet.Remove(addr) + } + + for _, address := range addresses { + if err := c.nodeService.DeleteNode(ctx, address); err != nil { + return errors.Wrapf(err, "deleting node at address: %s", address) + } + } + + if err := c.poller.RemoveAddresses(ctx, addresses...); err != nil { + return NewErrInternal(err.Error()) + } + + // No need to send Directives if nothing has ultimately changed. + if len(workerSet) == 0 { + return nil + } + + // Convert the slice of addresses into a slice of addressMethod containing + // the appropriate method. + addressMethods := applyAddressMethod(workerSet.SortedSlice(), dax.DirectiveMethodDiff) + + // For the addresses which are being removed, set their method to "reset". + for i := range addressMethods { + for j := range addresses { + if addressMethods[i].address == addresses[j] { + addressMethods[i].method = dax.DirectiveMethodReset + } + } + } + + // Get the current job assignments for these workers and send them to the + // nodes as Directives. + if err := c.sendDirectives(ctx, addressMethods...); err != nil { + return NewErrDirectiveSendFailure(err.Error()) + } + + return nil +} + +// Nodes returns the list of assigned nodes responsible for the jobs included in +// the given role. If createMissing is true, the Controller will create new jobs +// for any of which it isn't currently aware. +func (c *Controller) Nodes(ctx context.Context, role dax.Role, createMissing bool) ([]dax.AssignedNode, error) { + nodes := []dax.AssignedNode{} + var err error + + switch v := role.(type) { + case *dax.ComputeRole: + nodes, err = c.nodesCompute(ctx, v, createMissing) + if err != nil { + return nil, errors.Wrap(err, "getting compute nodes") + } + + case *dax.TranslateRole: + nodes, err = c.nodesTranslate(ctx, v, createMissing) + if err != nil { + return nil, errors.Wrap(err, "getting translate nodes") + } + } + + return nodes, nil +} + +// nodesTranslate is like nodesCompute. See the comments there. +func (c *Controller) nodesTranslate(ctx context.Context, role *dax.TranslateRole, createMissing bool) ([]dax.AssignedNode, error) { + // Try calling c.nodesTranslate as a read first. If we don't have to actually + // create any missing partitions, then we won't have to obtain a write lock. + translateNodes, retryAsWrite, err := c.nodesTranslateReadOrWrite(ctx, role, createMissing, false) + if err != nil { + return nil, errors.Wrap(err, "getting translate nodes read or write") + } + + if retryAsWrite { + translateNodes, _, err = c.nodesTranslateReadOrWrite(ctx, role, createMissing, retryAsWrite) + if err != nil { + return nil, errors.Wrap(err, "getting translate nodes read or write retry") + } + } + + return translateNodes, nil +} + +func (c *Controller) nodesTranslateReadOrWrite(ctx context.Context, role *dax.TranslateRole, createMissing bool, asWrite bool) ([]dax.AssignedNode, bool, error) { + if asWrite { + c.mu.Lock() + defer c.mu.Unlock() + } else { + c.mu.RLock() + defer c.mu.RUnlock() + } + + nodes := []dax.AssignedNode{} + + bal := c.TranslateBalancer + + //inJobs := NewStringSet() + inJobs := dax.NewSet[dax.Job]() + for _, p := range role.Partitions { + partitionString := partition(role.TableKey, p).String() + inJobs.Add(dax.Job(partitionString)) + } + + workers, err := bal.WorkersForJobs(ctx, inJobs.Sorted()) + if err != nil { + return nil, false, errors.Wrap(err, "getting workers for jobs") + } + + if createMissing { + // If any provided jobs were not returned in the WorkersForJobs + // request, then create those. + outJobs := dax.NewSet[dax.Job]() + for _, worker := range workers { + for _, job := range worker.Jobs { + outJobs.Add(job) + } + } + + missed := inJobs.Minus(outJobs).Sorted() + + if len(missed) > 0 { + // If we are currently under a read lock, and we get to this point, + // it means that we have partitions which need to be assigned (and + // directives sent) to workers. In that case, we need to abort this + // method run and notify the caller to rety as a write. + if !asWrite { + return nil, true, nil + } + + sort.Slice(missed, func(i, j int) bool { return missed[i] < missed[j] }) + + workerSet := NewAddressSet() + for _, job := range missed { + j, err := decodePartition(job) + if err != nil { + return nil, false, NewErrInternal(err.Error()) + } + diffs, err := bal.AddJob(ctx, j) + if err != nil { + return nil, false, errors.Wrap(err, "adding job") + } + for _, diff := range diffs { + workerSet.Add(dax.Address(diff.WorkerID)) + } + + // Initialize the partition version to 0. + qtid := j.table().QualifiedTableID() + if err := c.versionStore.AddPartitions(ctx, qtid, dax.NewPartition(j.partitionNum(), 0)); err != nil { + return nil, false, NewErrInternal(err.Error()) + } + } + + // Convert the slice of addresses into a slice of addressMethod containing + // the appropriate method. + addressMethods := applyAddressMethod(workerSet.SortedSlice(), dax.DirectiveMethodDiff) + + if err := c.sendDirectives(ctx, addressMethods...); err != nil { + return nil, false, NewErrDirectiveSendFailure(err.Error()) + } + + // Re-run WorkersForJobs. + workers, err = bal.WorkersForJobs(ctx, inJobs.Sorted()) + if err != nil { + return nil, false, errors.Wrap(err, "getting workers for jobs") + } + } + } + + for _, worker := range workers { + // covert worker.Jobs []string to map[string][]Partition + translateMap := make(map[dax.TableKey]dax.Partitions) + for _, job := range worker.Jobs { + j, err := decodePartition(job) + if err != nil { + return nil, false, NewErrInternal(err.Error()) + } + + // Get the partition version from the local versionStore. + tkey := j.table() + qtid := tkey.QualifiedTableID() + partitionVersion, found, err := c.versionStore.PartitionVersion(ctx, qtid, j.partitionNum()) + if err != nil { + return nil, false, err + } else if !found { + return nil, false, NewErrInternal("partition version not found in cache") + } + + translateMap[tkey] = append(translateMap[tkey], + dax.NewPartition(j.partitionNum(), partitionVersion), + ) + } + + for table, partitions := range translateMap { + // Sort the partitions int slice before returning it. + sort.Sort(partitions) + + nodes = append(nodes, dax.AssignedNode{ + Address: dax.Address(worker.ID), + Role: &dax.TranslateRole{ + TableKey: table, + Partitions: partitions, + }, + }) + } + } + + return nodes, false, nil +} + +// nodesCompute tries to get the list of compute nodes under a read lock. If the +// call into c.nodesComputeReadOrWrite() comes back with `retryAsWrite = true`, +// then it gets called again but with a write lock so that sendDirective can +// happen, and the receiving compute node can get apply the latest schema, +// without encountering race conditions. +// +// Really, we shouldn't have to rely on the directive being applied within a +// mu.Lock(). Instead, if a client (or the IDK) tries to perform some action on +// a compute node (for example, ingesting data to an index/field), if that +// action fails because the schema on the compute node is not in sync, or if the +// compute node is completely unavailable, the client should ask mds for updated +// node information and keep retrying. Basically, what I'm saying is that a lot +// of the mu.Lock()s in this file can be changed back to mu.RLock()s, and the +// SendDirective() can happen asyncronously without worring about race +// conditions. +// +// The race condition happened when concurrent requests to Nodes() occurred and +// the order of events was: +// - req1 wants node for [idx, 0] +// - req2 wants node for [idx, 0] +// - (req1) [idx, 0] registered in controller to node A +// - directive sent to node A to create index [idx] +// +// - (req2) receives: node A +// - (req2) tries to ingest data to node A [idx, 0] +// **** RACE: [idx] does not exist because (req1) directive step is not compete +func (c *Controller) nodesCompute(ctx context.Context, role *dax.ComputeRole, createMissing bool) ([]dax.AssignedNode, error) { + if len(role.Shards) == 0 { + return c.nodesForTableKey(ctx, role.TableKey) + } + // Try calling c.nodesCompute as a read first. If we don't have to actually + // create any missing shards, then we won't have to obtain a write lock. + computeNodes, retryAsWrite, err := c.nodesComputeReadOrWrite(ctx, role, createMissing, false) + if err != nil { + return nil, errors.Wrap(err, "getting compute nodes read or write") + } + + if retryAsWrite { + computeNodes, _, err = c.nodesComputeReadOrWrite(ctx, role, createMissing, retryAsWrite) + if err != nil { + return nil, errors.Wrap(err, "getting compute nodes read or write retry") + } + } + + return computeNodes, nil +} + +func (c *Controller) nodesForTableKey(ctx context.Context, tk dax.TableKey) ([]dax.AssignedNode, error) { + bal := c.ComputeBalancer + workers, err := bal.WorkersForJobPrefix(ctx, string(tk)) + if err != nil { + return nil, errors.Wrapf(err, "getting workers for table: '%s'", tk) + } + + return c.workersToAssignedNodes(ctx, workers) + +} + +// nodesComputeReadOrWrite contains the logic for the c.nodesCompute() method, +// but it supports being called with either a read or write lock. +func (c *Controller) nodesComputeReadOrWrite(ctx context.Context, role *dax.ComputeRole, createMissing bool, asWrite bool) ([]dax.AssignedNode, bool, error) { + if asWrite { + c.mu.Lock() + defer c.mu.Unlock() + } else { + c.mu.RLock() + defer c.mu.RUnlock() + } + + bal := c.ComputeBalancer + + inJobs := dax.NewSet[dax.Job]() + for _, s := range role.Shards { + shardString := shard(role.TableKey, s).String() + inJobs.Add(dax.Job(shardString)) + } + + workers, err := bal.WorkersForJobs(ctx, inJobs.Sorted()) + if err != nil { + return nil, false, errors.Wrap(err, "getting workers for jobs") + } + + // figure out if any jobs in the role have no workers assigned + outJobs := dax.NewSet[dax.Job]() + for _, worker := range workers { + for _, job := range worker.Jobs { + outJobs.Add(job) + } + } + + missed := inJobs.Minus(outJobs).Sorted() + if !createMissing && len(missed) > 0 { + return nil, false, NewErrUnassignedJobs(missed) + } + + if createMissing { + // If any provided jobs were not returned in the WorkersForJobs + // request, then create those. + if len(missed) > 0 { + // If we are currently under a read lock, and we get to this point, + // it means that we have shards which need to be assigned (and + // directives sent) to workers. In that case, we need to abort this + // method run and notify the caller to rety as a write. + if !asWrite { + return nil, true, nil + } + + sort.Slice(missed, func(i, j int) bool { return missed[i] < missed[j] }) + + workerSet := NewAddressSet() + for _, job := range missed { + j, err := decodeShard(job) + if err != nil { + return nil, false, NewErrInternal(err.Error()) + } + diffs, err := bal.AddJob(ctx, j) + if err != nil { + return nil, false, errors.Wrap(err, "adding job") + } + for _, diff := range diffs { + workerSet.Add(dax.Address(diff.WorkerID)) + } + + // Initialize the shard version to 0. + qtid := j.table().QualifiedTableID() + if err := c.versionStore.AddShards(ctx, qtid, dax.NewShard(j.shardNum(), 0)); err != nil { + return nil, false, NewErrInternal(err.Error()) + } + } + + // Convert the slice of addresses into a slice of addressMethod containing + // the appropriate method. + addressMethods := applyAddressMethod(workerSet.SortedSlice(), dax.DirectiveMethodDiff) + + if err := c.sendDirectives(ctx, addressMethods...); err != nil { + return nil, false, NewErrDirectiveSendFailure(err.Error()) + } + + // Re-run WorkersForJobs. + workers, err = bal.WorkersForJobs(ctx, inJobs.Sorted()) + if err != nil { + return nil, false, errors.Wrap(err, "getting workers for jobs") + } + } + } + + nodes, err := c.workersToAssignedNodes(ctx, workers) + return nodes, false, errors.Wrap(err, "converting to assigned nodes") +} + +func (c *Controller) workersToAssignedNodes(ctx context.Context, workers []dax.WorkerInfo) ([]dax.AssignedNode, error) { + nodes := []dax.AssignedNode{} + for _, worker := range workers { + // convert worker.Jobs []string to map[TableName][]Shard + computeMap := make(map[dax.TableKey]dax.Shards) + for _, job := range worker.Jobs { + j, err := decodeShard(job) + if err != nil { + return nil, NewErrInternal(err.Error()) + } + + // Get the shard version from the local versionStore. + tkey := j.table() + qtid := tkey.QualifiedTableID() + shardVersion, found, err := c.versionStore.ShardVersion(ctx, qtid, j.shardNum()) + if err != nil { + return nil, err + } else if !found { + return nil, NewErrInternal("shard version not found in cache") + } + + computeMap[tkey] = append(computeMap[tkey], + dax.NewShard(j.shardNum(), shardVersion), + ) + } + + for table, shards := range computeMap { + // Sort the shards uint64 slice before returning it. + sort.Sort(shards) + + nodes = append(nodes, dax.AssignedNode{ + Address: dax.Address(worker.ID), + Role: &dax.ComputeRole{ + TableKey: table, + Shards: shards, + }, + }) + } + } + return nodes, nil +} + +// CreateTable adds a table to the versionStore and schemar, and then sends directives +// to all affected nodes based on the change. +func (c *Controller) CreateTable(ctx context.Context, qtbl *dax.QualifiedTable) error { + c.mu.Lock() + defer c.mu.Unlock() + + qtid := qtbl.QualifiedID() + + // Add the table to the versionStore. + if err := c.versionStore.AddTable(ctx, qtid); err != nil { + return errors.Wrapf(err, "adding table: %s", qtid) + } + + // Add fields which have string keys to the local versionStore. + fieldVersions := make(dax.FieldVersions, 0) + for _, field := range qtbl.Fields { + if !field.StringKeys() { + continue + } + + fieldVersions = append(fieldVersions, dax.FieldVersion{ + Name: field.Name, + Version: 0, + }) + } + if len(fieldVersions) > 0 { + if err := c.versionStore.AddFields(ctx, qtid, fieldVersions...); err != nil { + return errors.Wrapf(err, "adding fields: %s, %v", qtid, fieldVersions) + } + } + + // If the table is keyed, add partitions to the balancer. + if qtbl.StringKeys() { + // workerSet maintains the set of workers which have a job assignment change + // and therefore need to be sent an updated Directive. + workerSet := NewAddressSet() + + // Generate the list of partitions to be added. + partitions := make(dax.Partitions, qtbl.PartitionN) + for partitionNum := 0; partitionNum < qtbl.PartitionN; partitionNum++ { + partitions[partitionNum] = dax.NewPartition(dax.PartitionNum(partitionNum), 0) + } + + // Add partitions to versionStore. Version is intentionally set to 0 + // here as this is the initial instance of the partition. + if err := c.versionStore.AddPartitions(ctx, qtid, partitions...); err != nil { + return NewErrInternal(err.Error()) + } + + for _, p := range partitions { + // We don't currently use the returned diff, other than to determine + // which worker was affected, because we send the full Directive + // every time. + diffs, err := c.TranslateBalancer.AddJob(ctx, partition(qtbl.Key(), p)) + if err != nil { + return errors.Wrap(err, "adding job") + } + for _, diff := range diffs { + workerSet.Add(dax.Address(diff.WorkerID)) + } + } + + // Convert the slice of addresses into a slice of addressMethod containing + // the appropriate method. + addressMethods := applyAddressMethod(workerSet.SortedSlice(), dax.DirectiveMethodDiff) + + if err := c.sendDirectives(ctx, addressMethods...); err != nil { + return NewErrDirectiveSendFailure(err.Error()) + } + } + + // This is more FieldVersion hackery. Even if the table is not keyed, we + // still want to manage partition 0 for the table in case any of the table's + // fields contain string keys (we use partition 0 for field string keys for + // now; in the future we should distribute/balance the field key translation + // like we do shards and partitions). + if !qtbl.StringKeys() { + // workerSet maintains the set of workers which have a job assignment change + // and therefore need to be sent an updated Directive. + workerSet := NewAddressSet() + + p := dax.NewPartition(0, 0) + + // Add partition 0 to versionStore. Version is intentionally set to 0 + // here as this is the initial instance of the partition. + if err := c.versionStore.AddPartitions(ctx, qtid, p); err != nil { + return NewErrInternal(err.Error()) + } + + // We don't currently use the returned diff, other than to determine + // which worker was affected, because we send the full Directive + // every time. + diffs, err := c.TranslateBalancer.AddJob(ctx, partition(qtbl.Key(), p)) + if err != nil { + return errors.Wrap(err, "adding job") + } + for _, diff := range diffs { + workerSet.Add(dax.Address(diff.WorkerID)) + } + + // Convert the slice of addresses into a slice of addressMethod containing + // the appropriate method. + addressMethods := applyAddressMethod(workerSet.SortedSlice(), dax.DirectiveMethodDiff) + + if err := c.sendDirectives(ctx, addressMethods...); err != nil { + return NewErrDirectiveSendFailure(err.Error()) + } + } + + return nil +} + +// DropTable removes a table from the schema and sends directives to all affected +// nodes based on the change. +func (c *Controller) DropTable(ctx context.Context, qtid dax.QualifiedTableID) error { + c.mu.Lock() + defer c.mu.Unlock() + + // Get the table from the schemar. + if _, err := c.Schemar.Table(ctx, qtid); err != nil { + return errors.Wrapf(err, "table not in schemar: %s", qtid) + } + + // Remove the table from the versionStore. + // Since the schemar should be the system of record for the existence of a + // table, if the versionStore is not aware of the table, we just log it and + // continue. + shards, partitions, err := c.versionStore.RemoveTable(ctx, qtid) + if err != nil { + return errors.Wrapf(err, "removing table: %s", qtid) + } + + // workerSet maintains the set of workers which have a job assignment change + // and therefore need to be sent an updated Directive. + workerSet := NewAddressSet() + + // Remove shards. + for _, s := range shards { + diffs, err := c.ComputeBalancer.RemoveJob(ctx, shard(qtid.Key(), s)) + if err != nil { + return errors.Wrap(err, "removing job") + } + for _, diff := range diffs { + workerSet.Add(dax.Address(diff.WorkerID)) + } + } + + // Remove partitions. + for _, p := range partitions { + diffs, err := c.TranslateBalancer.RemoveJob(ctx, partition(qtid.Key(), p)) + if err != nil { + return errors.Wrap(err, "removing job") + } + for _, diff := range diffs { + workerSet.Add(dax.Address(diff.WorkerID)) + } + } + + // Convert the slice of addresses into a slice of addressMethod containing + // the appropriate method. + addressMethods := applyAddressMethod(workerSet.SortedSlice(), dax.DirectiveMethodDiff) + + if err := c.sendDirectives(ctx, addressMethods...); err != nil { + return NewErrDirectiveSendFailure(err.Error()) + } + + return nil +} + +// Table returns a table by quaified table id. +func (c *Controller) Table(ctx context.Context, qtid dax.QualifiedTableID) (*dax.QualifiedTable, error) { + c.mu.Lock() + defer c.mu.Unlock() + + // Get the table from the schemar. + return c.Schemar.Table(ctx, qtid) +} + +// Tables returns a list of tables by name. +func (c *Controller) Tables(ctx context.Context, qual dax.TableQualifier, ids ...dax.TableID) ([]*dax.QualifiedTable, error) { + c.mu.Lock() + defer c.mu.Unlock() + + // Get the tables from the schemar. + return c.Schemar.Tables(ctx, qual, ids...) +} + +// AddShards registers the table/shard combinations with the controller and +// sends the necessary directive. +func (c *Controller) AddShards(ctx context.Context, qtid dax.QualifiedTableID, shards ...dax.Shard) error { + c.mu.Lock() + defer c.mu.Unlock() + + // Add shards to versionStore. + if err := c.versionStore.AddShards(ctx, qtid, shards...); err != nil { + return errors.Wrapf(err, "adding shards: %s, %v", qtid, shards) + } + + // workerSet maintains the set of workers which have a job assignment change + // and therefore need to be sent an updated Directive. + workerSet := NewAddressSet() + + for _, s := range shards { + // We don't currently use the returned diff, other than to determine + // which worker was affected, because we send the full Directive every + // time. + diffs, err := c.ComputeBalancer.AddJob(ctx, shard(qtid.Key(), s)) + if err != nil { + return errors.Wrap(err, "adding job") + } + for _, diff := range diffs { + workerSet.Add(dax.Address(diff.WorkerID)) + } + + // Initialize the shard version to 0. + if err := c.versionStore.AddShards(ctx, qtid, dax.NewShard(s.Num, 0)); err != nil { + return NewErrInternal(err.Error()) + } + } + + // Convert the slice of addresses into a slice of addressMethod containing + // the appropriate method. + addressMethods := applyAddressMethod(workerSet.SortedSlice(), dax.DirectiveMethodDiff) + + if err := c.sendDirectives(ctx, addressMethods...); err != nil { + return NewErrDirectiveSendFailure(err.Error()) + } + + return nil +} + +// RemoveShards deregisters the table/shard combinations with the controller and +// sends the necessary directives. +func (c *Controller) RemoveShards(ctx context.Context, qtid dax.QualifiedTableID, shards ...dax.Shard) error { + c.mu.Lock() + defer c.mu.Unlock() + + // workerSet maintains the set of workers which have a job assignment change + // and therefore need to be sent an updated Directive. + workerSet := NewAddressSet() + + for _, s := range shards { + // We don't currently use the returned diff, other than to determine + // which worker was affected, because we send the full Directive every + // time. + diffs, err := c.ComputeBalancer.RemoveJob(ctx, shard(qtid.Key(), s)) + if err != nil { + return errors.Wrap(err, "removing job") + } + for _, diff := range diffs { + workerSet.Add(dax.Address(diff.WorkerID)) + } + } + + // Convert the slice of addresses into a slice of addressMethod containing + // the appropriate method. + addressMethods := applyAddressMethod(workerSet.SortedSlice(), dax.DirectiveMethodDiff) + + if err := c.sendDirectives(ctx, addressMethods...); err != nil { + return NewErrDirectiveSendFailure(err.Error()) + } + + return nil +} + +// sendDirectives sends a directive (based on the current balancer state) to +// each of the nodes provided. +func (c *Controller) sendDirectives(ctx context.Context, addrs ...addressMethod) error { + // If nodes is empty, return early. + if len(addrs) == 0 { + return nil + } + + directives, err := c.buildDirectives(ctx, addrs, c.versionStore) + if err != nil { + return errors.Wrap(err, "building directives") + } + + errs := make([]error, len(directives)) + var eg errgroup.Group + for i, dir := range directives { + i := i + dir := dir + eg.Go(func() error { + errs[i] = c.Director.SendDirective(ctx, dir) + if dir.IsEmpty() { + errs[i] = nil + } + return errs[i] + }) + } + + err = eg.Wait() + if err != nil { + errCount := 0 + for _, err := range errs { + if err != nil { + errCount++ + } + } + if doWeCare(directives, errs) { + // TODO: in this case, we should probably remove nodes + // that didn't work and retry. + return errors.Errorf("all directives errored: %+v", errs) + } + } + + return nil +} + +func doWeCare(directives []*dax.Directive, errors []error) bool { + for i, directive := range directives { + err := errors[i] + if err != nil && !directive.IsEmpty() { + return true + } + } + return false +} + +// addressMethod is used when building a Directive to specify which +// DirectiveMethod should be applied for the given Address. +type addressMethod struct { + address dax.Address + method dax.DirectiveMethod +} + +// applyAddressMethod converts the slice of addrs into a slice of addressMethod +// containing the given method. +func applyAddressMethod(addrs []dax.Address, method dax.DirectiveMethod) []addressMethod { + ams := make([]addressMethod, len(addrs)) + for i := range addrs { + ams[i] = addressMethod{ + address: addrs[i], + method: method, + } + } + + return ams +} + +// buildDirectives builds a list of directives for the given addrs (i.e. nodes) +// using information (i.e. current state) from the balancers. +func (c *Controller) buildDirectives(ctx context.Context, addrs []addressMethod, versionStore dax.VersionStore) ([]*dax.Directive, error) { + directives := make([]*dax.Directive, len(addrs)) + + for i, addressMethod := range addrs { + dVersion, err := c.directiveVersion.Increment(ctx, 1) + if err != nil { + return nil, errors.Wrap(err, "incrementing directive version") + } + + d := &dax.Directive{ + Address: addressMethod.address, + Method: addressMethod.method, + Tables: []*dax.QualifiedTable{}, + ComputeRoles: []dax.ComputeRole{}, + TranslateRoles: []dax.TranslateRole{}, + Version: dVersion, + } + + // computeMap maps a table to a list of shards for that table. We need + // to aggregate them here because the list of jobs from WorkerState() + // can contain a mixture of table/shards. + computeMap := make(map[dax.TableKey][]dax.Shard) + + // translateMap maps a table to a list of partitions for that table. We + // need to aggregate them here because the list of jobs from + // WorkerState() can contain a mixture of table/partitions. + translateMap := make(map[dax.TableKey]dax.Partitions) + + // tableSet maintains the set of tables which have a job assignment + // change and therefore need to be included in the Directive schema. + tableSet := NewTableSet() + + // ownsPartition0 is the list of tables for which this node owns partition 0. + // This is used to determine FieldVersion responsiblity. + ownsPartition0 := make(map[dax.TableKey]struct{}, 0) + + for _, rt := range supportedRoleTypes { + bal, err := c.balancerForRole(rt) + if err != nil { + return nil, errors.Wrap(err, "getting balancer") + } + w, err := bal.WorkerState(ctx, dax.Worker(addressMethod.address.String())) + if err != nil { + return nil, errors.Wrapf(err, "getting worker state: %s", addressMethod.address) + } + + switch rt { + case dax.RoleTypeCompute: + for _, job := range w.Jobs { + j, err := decodeShard(job) + if err != nil { + return nil, errors.Wrapf(err, "decoding shard job: %s", job) + } + + // The Shard object decoded from the balancer doesn't + // contain a valid version (because the balancer + // intentionally does not store version information). Here, + // we get the current shard version from the controller's + // cache (i.e. versionStore) and inject that into the Shard + // sent in the directive. + tkey := j.table() + qtid := tkey.QualifiedTableID() + shardVersion, found, err := versionStore.ShardVersion(ctx, qtid, j.shardNum()) + if err != nil { + return nil, errors.Wrapf(err, "getting shard version: %s, %d", qtid, j.shardNum()) + } else if !found { + return nil, NewErrInternal("shard version not found in cache") + } + + computeMap[tkey] = append(computeMap[tkey], + dax.NewShard(j.shardNum(), shardVersion), + ) + tableSet.Add(tkey) + } + case dax.RoleTypeTranslate: + for _, job := range w.Jobs { + j, err := decodePartition(job) + if err != nil { + return nil, errors.Wrapf(err, "decoding partition job: %s", job) + } + + // This check is related to the FieldVersion logic below. + // Basically, we need to determine if this node is + // responsible for partition 0 for any table(s), and if so, + // include FieldVersion in the directive for the node. + if j.partitionNum() == 0 { + ownsPartition0[j.table()] = struct{}{} + } + + // The Partition object decoded from the balancer doesn't + // contain a valid version (because the balancer + // intentionally does not store version information). Here, + // we get the current partition version from the + // controller's cache (i.e. versionStore) and inject that + // into the Partition sent in the directive. + tkey := j.table() + qtid := tkey.QualifiedTableID() + partitionVersion, found, err := versionStore.PartitionVersion(ctx, qtid, j.partitionNum()) + if err != nil { + return nil, errors.Wrapf(err, "getting partition version: %s, %d", qtid, j.partitionNum()) + } else if !found { + return nil, NewErrInternal("partition version not found in cache") + } + + translateMap[tkey] = append(translateMap[tkey], + dax.NewPartition(j.partitionNum(), partitionVersion), + ) + tableSet.Add(tkey) + } + } + } + + // Convert the computeMap into a list of ComputeRole. + for k, v := range computeMap { + // Because these were encoded as strings in the balancer and may be + // out of order numerically, sort them as integers. + //sort.Slice(v, func(i, j int) bool { return v[i] < v[j] }) + sort.Sort(dax.Shards(v)) + + d.ComputeRoles = append(d.ComputeRoles, dax.ComputeRole{ + TableKey: k, + Shards: v, + }) + } + + // Convert the translateMap into a list of TranslateRole. + for k, v := range translateMap { + // Because these were encoded as strings in the balancer and may be + // out of order numerically, sort them as integers. + sort.Sort(v) + + d.TranslateRoles = append(d.TranslateRoles, dax.TranslateRole{ + TableKey: k, + Partitions: v, + }) + } + + // Add field-specific TranslateRoles to the node which is responsible + // for partition 0. This is a bit clunkly; ideally we would handle this + // the same way we handle shards and partitions, by maintaining a + // distinct balancer for FieldVersions. But because the query side isn't + // currently set up to look for field translation anywhere but on the + // local node (or in the case of MDS, on partition 0), we're keeping + // everything that way for now. + for tkey := range ownsPartition0 { + qtid := tkey.QualifiedTableID() + table, err := c.Schemar.Table(ctx, qtid) + if err != nil { + return nil, errors.Wrapf(err, "getting table: %s", tkey) + } + + fieldVersions := make(dax.FieldVersions, 0) + for _, field := range table.Fields { + if !field.StringKeys() { + continue + } + + // Skip the primary key field; it uses table translation. + if field.IsPrimaryKey() { + continue + } + + fieldVersion, found, err := versionStore.FieldVersion(ctx, qtid, field.Name) + if err != nil { + return nil, errors.Wrapf(err, "getting field version: %s, %s", qtid, field) + } else if !found { + return nil, NewErrInternal("field version not found in cache") + } + + fieldVersions = append(fieldVersions, dax.FieldVersion{ + Name: field.Name, + Version: fieldVersion, + }) + } + + if len(fieldVersions) == 0 { + continue + } + + d.TranslateRoles = append(d.TranslateRoles, dax.TranslateRole{ + TableKey: tkey, + Fields: fieldVersions, + }) + + tableSet.Add(tkey) + } + /////////////// end of FieldVersion logic ////////////////////// + + if len(tableSet) > 0 { + dTables := make([]*dax.QualifiedTable, 0) + for qual, tblIDs := range tableSet.QualifiedSortedSlice() { + qtbls, err := c.Schemar.Tables(ctx, qual, tblIDs...) + if err != nil { + return nil, errors.Wrapf(err, "getting directive tables for qual: %s", qual) + } + dTables = append(dTables, qtbls...) + } + d.Tables = dTables + } + + // Sort ComputeRoles by table. + sort.Slice(d.ComputeRoles, func(i, j int) bool { return d.ComputeRoles[i].TableKey < d.ComputeRoles[j].TableKey }) + + // Sort TranslateRoles by table. + sort.Slice(d.TranslateRoles, func(i, j int) bool { return d.TranslateRoles[i].TableKey < d.TranslateRoles[j].TableKey }) + + directives[i] = d + } + + return directives, nil +} + +func (c *Controller) SetPoller(poller dax.AddressManager) { + c.poller = poller +} + +// InitializePoller sends the list of known nodes (to be polled) to the poller. +// This is useful in the case where MDS has restarted (or has been replaced) and +// its poller is emtpy (i.e. it doesn't know about any nodes). +func (c *Controller) InitializePoller(ctx context.Context) error { + nodes, err := c.nodeService.Nodes(context.Background()) + if err != nil { + return errors.Wrap(err, "initializing poller") + } + + for _, node := range nodes { + if err := c.poller.AddAddresses(ctx, node.Address); err != nil { + return errors.Wrapf(err, "adding address to poller: %s", node.Address) + } + } + + return nil +} + +// SnapshotTable snapshots a table. +func (c *Controller) SnapshotTable(ctx context.Context, qtid dax.QualifiedTableID) error { + shards, ok, err := c.versionStore.Shards(ctx, qtid) + if err != nil { + return errors.Wrap(err, "getting shards from version store") + } else if !ok { + return errors.New(errors.ErrUncoded, "got false back from versionStore.Shards") + } + + for _, shard := range shards { + if err := c.SnapshotShardData(ctx, qtid, shard.Num); err != nil { + return errors.Wrapf(err, "snapshotting shard data: qtid: %s, shard: %d", qtid, shard.Num) + } + } + + partitions, ok, err := c.versionStore.Partitions(ctx, qtid) + if err != nil { + return errors.Wrap(err, "getting partitions from version store") + } else if !ok { + return errors.New(errors.ErrUncoded, "got false back from versionStore.Partitions") + } + + for _, part := range partitions { + if err := c.SnapshotTableKeys(ctx, qtid, part.Num); err != nil { + return errors.Wrapf(err, "snapshotting table keys: qtid: %s, partition: %d", qtid, part.Num) + } + } + + fields, ok, err := c.versionStore.Fields(ctx, qtid) + if err != nil { + return errors.Wrap(err, "getting fields from version store") + } else if !ok { + return errors.New(errors.ErrUncoded, "got false back from versionStore.Fields") + } + + for _, fld := range fields { + if fld.Name != "_id" { + if err := c.SnapshotFieldKeys(ctx, qtid, fld.Name); err != nil { + return errors.Wrapf(err, "snapshotting field keys: qtid: %s, field: %s", qtid, fld.Name) + } + } + } + return nil +} + +// SnapshotShardData forces the compute node responsible for the given shard to +// snapshot that shard, then increment its shard version for logs written to the +// WriteLogger. +func (c *Controller) SnapshotShardData(ctx context.Context, qtid dax.QualifiedTableID, shardNum dax.ShardNum) error { + // Confirm table/shard is being tracked; get the current shard. + fromShardVersion, ok, err := c.versionStore.ShardVersion(ctx, qtid, shardNum) + if err != nil { + return errors.Wrapf(err, "getting shard version: %s, %d", qtid, shardNum) + } else if !ok { + return NewErrInternal( + fmt.Sprintf("shard to snapshot not found: %s, %d", qtid, shardNum), + ) + } + toShardVersion := fromShardVersion + 1 + + // Get the node responsible for the shard. + bal := c.ComputeBalancer + + job := shard(qtid.Key(), dax.NewShard(shardNum, -1)) + + workers, err := bal.WorkersForJobs(ctx, []dax.Job{dax.Job(job.String())}) + if err != nil { + return errors.Wrapf(err, "getting workers for jobs: %s", job) + } + if len(workers) == 0 { + c.logger.Printf("no worker found for shard: %s, %d", qtid, shardNum) + return nil + } + + addr := dax.Address(workers[0].ID) + + // Make a copy of the controller's versionStore, and update the current + // shard so that the directive sent along with the SnapshotRequest reflects + // the state that we want after a successful snapshot. + versionStoreCopy, err := c.versionStore.Copy(ctx) + if err != nil { + return errors.Wrap(err, "copying version store") + } + if err := versionStoreCopy.AddShards(ctx, qtid, + dax.NewShard(shardNum, toShardVersion), + ); err != nil { + return NewErrInternal(err.Error()) + } + + // Convert the address into a slice of addressMethod containing the + // appropriate method. + addressMethods := applyAddressMethod([]dax.Address{addr}, dax.DirectiveMethodSnapshot) + + var toDirective dax.Directive + if directives, err := c.buildDirectives(ctx, addressMethods, versionStoreCopy); err != nil { + return NewErrInternal(err.Error()) + } else if ld := len(directives); ld != 1 { + msg := fmt.Sprintf("buildDirectives returned invalid number of directives: %d", ld) + return NewErrInternal(msg) + } else { + toDirective = *directives[0] + } + + // Send the node a snapshot request. + req := &dax.SnapshotShardDataRequest{ + Address: addr, + TableKey: qtid.Key(), + ShardNum: shardNum, + FromVersion: fromShardVersion, + ToVersion: toShardVersion, + Directive: toDirective, + } + + if err := c.Director.SendSnapshotShardDataRequest(ctx, req); err != nil { + return NewErrInternal(err.Error()) + } + + // A successful request means the shard version can be incremented. + if err := c.versionStore.AddShards(ctx, qtid, + dax.NewShard(shardNum, toShardVersion), + ); err != nil { + return NewErrInternal(err.Error()) + } + + return nil +} + +// SnapshotTableKeys forces the translate node responsible for the given +// partition to snapshot the table keys for that partition, then increment its +// version for logs written to the WriteLogger. +func (c *Controller) SnapshotTableKeys(ctx context.Context, qtid dax.QualifiedTableID, partitionNum dax.PartitionNum) error { + // Confirm table/shard is being tracked; get the current shard. + fromPartitionVersion, found, err := c.versionStore.PartitionVersion(ctx, qtid, partitionNum) + if err != nil { + return errors.Wrapf(err, "getting partition version: %s, %d", qtid, partitionNum) + } else if !found { + return NewErrInternal( + fmt.Sprintf("partition to snapshot not found: %s, %d", qtid, partitionNum), + ) + } + toPartitionVersion := fromPartitionVersion + 1 + + // Get the node responsible for the partition. + bal := c.TranslateBalancer + + job := partition(qtid.Key(), dax.NewPartition(partitionNum, -1)) + + workers, err := bal.WorkersForJobs(ctx, []dax.Job{dax.Job(job.String())}) + if err != nil { + return errors.Wrapf(err, "getting workers for jobs: %s", job) + } + if len(workers) == 0 { + c.logger.Printf("no worker found for partition: %s, %d", qtid, partitionNum) + return nil + } + + addr := dax.Address(workers[0].ID) + + // Make a copy of the controller's versionStore, and update the current + // partition so that the directive sent along with the SnapshotRequest + // reflects the state that we want after a successful snapshot. + versionStoreCopy, err := c.versionStore.Copy(ctx) + if err != nil { + return errors.Wrap(err, "copying version store") + } + if err := versionStoreCopy.AddPartitions(ctx, qtid, + dax.NewPartition(partitionNum, toPartitionVersion), + ); err != nil { + return NewErrInternal(err.Error()) + } + + // Convert the address into a slice of addressMethod containing the + // appropriate method. + addressMethods := applyAddressMethod([]dax.Address{addr}, dax.DirectiveMethodSnapshot) + + var toDirective dax.Directive + if directives, err := c.buildDirectives(ctx, addressMethods, versionStoreCopy); err != nil { + return NewErrInternal(err.Error()) + } else if ld := len(directives); ld != 1 { + msg := fmt.Sprintf("buildDirectives returned invalid number of directives: %d", ld) + return NewErrInternal(msg) + } else { + toDirective = *directives[0] + } + + // Send the node a snapshot request. + req := &dax.SnapshotTableKeysRequest{ + Address: addr, + TableKey: qtid.Key(), + PartitionNum: partitionNum, + FromVersion: fromPartitionVersion, + ToVersion: toPartitionVersion, + Directive: toDirective, + } + + if err := c.Director.SendSnapshotTableKeysRequest(ctx, req); err != nil { + return NewErrInternal(err.Error()) + } + + // A successful request means the partition version can be incremented. + if err := c.versionStore.AddPartitions(ctx, qtid, + dax.NewPartition(partitionNum, toPartitionVersion), + ); err != nil { + return NewErrInternal(err.Error()) + } + + return nil +} + +// SnapshotFieldKeys forces the translate node responsible for the given field +// to snapshot the keys for that field, then increment its version for logs +// written to the WriteLogger. +func (c *Controller) SnapshotFieldKeys(ctx context.Context, qtid dax.QualifiedTableID, field dax.FieldName) error { + // Confirm table/field is being tracked; get the current field. + fromFieldVersion, ok, err := c.versionStore.FieldVersion(ctx, qtid, field) + if err != nil { + return errors.Wrapf(err, "getting field version: %s, %s", qtid, field) + } else if !ok { + return NewErrInternal( + fmt.Sprintf("field to snapshot not found: %s, %s", qtid, field), + ) + } + toFieldVersion := fromFieldVersion + 1 + + // Get the node responsible for the field. + bal := c.TranslateBalancer + + // Field translation is currently handled by partition 0. + partitionNum := dax.PartitionNum(0) + job := partition(qtid.Key(), dax.NewPartition(partitionNum, -1)) + + workers, err := bal.WorkersForJobs(ctx, []dax.Job{dax.Job(job.String())}) + if err != nil { + return errors.Wrapf(err, "getting workers for jobs: %s", job) + } + if len(workers) == 0 { + c.logger.Printf("no worker found for partition: %s, %d", qtid, partitionNum) + return nil + } + + addr := dax.Address(workers[0].ID) + + // Make a copy of the controller's versionStore, and update the current + // field so that the directive sent along with the SnapshotRequest reflects + // the state that we want after a successful snapshot. + versionStoreCopy, err := c.versionStore.Copy(ctx) + if err != nil { + return errors.Wrap(err, "copying version store") + } + if err := versionStoreCopy.AddFields(ctx, qtid, + dax.NewFieldVersion(field, toFieldVersion), + ); err != nil { + return NewErrInternal(err.Error()) + } + + // Convert the address into a slice of addressMethod containing the + // appropriate method. + addressMethods := applyAddressMethod([]dax.Address{addr}, dax.DirectiveMethodSnapshot) + + var toDirective dax.Directive + if directives, err := c.buildDirectives(ctx, addressMethods, versionStoreCopy); err != nil { + return NewErrInternal(err.Error()) + } else if ld := len(directives); ld != 1 { + msg := fmt.Sprintf("buildDirectives returned invalid number of directives: %d", ld) + return NewErrInternal(msg) + } else { + toDirective = *directives[0] + } + + // Send the node a snapshot request. + req := &dax.SnapshotFieldKeysRequest{ + Address: addr, + TableKey: qtid.Key(), + Field: field, + FromVersion: fromFieldVersion, + ToVersion: toFieldVersion, + Directive: toDirective, + } + + if err := c.Director.SendSnapshotFieldKeysRequest(ctx, req); err != nil { + return NewErrInternal(err.Error()) + } + + // A successful request means the field version can be incremented. + if err := c.versionStore.AddFields(ctx, qtid, + dax.NewFieldVersion(field, toFieldVersion), + ); err != nil { + return NewErrInternal(err.Error()) + } + + return nil +} + +///////////// + +func (c *Controller) ComputeNodes(ctx context.Context, qtid dax.QualifiedTableID, shards dax.ShardNums, isWrite bool) ([]ComputeNode, error) { + inRole := &dax.ComputeRole{ + TableKey: qtid.Key(), + Shards: dax.NewShards(shards...), + } + + nodes, err := c.Nodes(ctx, inRole, isWrite) + if err != nil { + return nil, errors.Wrap(err, "getting compute nodes") + } + + computeNodes := make([]ComputeNode, 0) + + for _, node := range nodes { + role, ok := node.Role.(*dax.ComputeRole) + if !ok { + // TODO: this should be impossible, but still, we could use + // some API (HTTP?) error codes. + return nil, NewErrInternal("not a compute node") + } + + computeNodes = append(computeNodes, ComputeNode{ + Address: node.Address, + Table: role.TableKey, + Shards: role.Shards.Nums(), + }) + } + + return computeNodes, nil +} + +func (c *Controller) TranslateNodes(ctx context.Context, qtid dax.QualifiedTableID, partitions dax.PartitionNums, isWrite bool) ([]TranslateNode, error) { + inRole := &dax.TranslateRole{ + TableKey: qtid.Key(), + Partitions: dax.NewPartitions(partitions...), + } + + nodes, err := c.Nodes(ctx, inRole, isWrite) + if err != nil { + return nil, errors.Wrap(err, "getting translate nodes") + } + + translateNodes := make([]TranslateNode, 0) + + for _, node := range nodes { + role, ok := node.Role.(*dax.TranslateRole) + if !ok { + // TODO: this should be impossible, but still, we could use + // some API (HTTP?) error codes. + return nil, NewErrInternal("not a translate node") + } + + translateNodes = append(translateNodes, TranslateNode{ + Address: node.Address, + Table: role.TableKey, + Partitions: role.Partitions.Nums(), + }) + } + + return translateNodes, nil +} + +func (c *Controller) IngestPartition(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum) (dax.Address, error) { + c.mu.RLock() + defer c.mu.RUnlock() + + partitions := dax.PartitionNums{partition} + + nodes, err := c.TranslateNodes(ctx, qtid, partitions, true) + if err != nil { + return "", err + } + + if l := len(nodes); l == 0 { + return "", NewErrNoAvailableNode() + } else if l > 1 { + return "", + NewErrInternal( + fmt.Sprintf("unexpected number of nodes: %d", l)) + } + + node := nodes[0] + + // Verify that the node returned is actually responsible for the partition + // requested. + if node.Table != qtid.Key() { + return "", + NewErrInternal( + fmt.Sprintf("table returned (%s) does not match requested (%s)", node.Table, qtid)) + } else if l := len(node.Partitions); l != 1 { + return "", + NewErrInternal( + fmt.Sprintf("unexpected number of partitions returned: %d", l)) + } else if p := node.Partitions[0]; p != partition { + return "", + NewErrInternal( + fmt.Sprintf("partition returned (%d) does not match requested (%d)", p, partition)) + } + + return node.Address, nil +} + +//// + +func (c *Controller) CreateField(ctx context.Context, qtid dax.QualifiedTableID, fld *dax.Field) error { + c.mu.Lock() + defer c.mu.Unlock() + + // workerSet maintains the set of workers which have a job assignment change + // and therefore need to be sent an updated Directive. + workerSet := NewAddressSet() + + // If the field has string keys, add it to the local versionStore. + if fld.StringKeys() { + fieldVersion := dax.FieldVersion{ + Name: fld.Name, + Version: 0, + } + if err := c.versionStore.AddFields(ctx, qtid, fieldVersion); err != nil { + return errors.Wrapf(err, "adding fields: %s, %s", qtid, fieldVersion) + } + } + + // Get the worker responsible for partition 0, which handles field key + // translation. Be sure to get the current version. + if v, found, err := c.versionStore.PartitionVersion(ctx, qtid, 0); err != nil { + return errors.Wrapf(err, "getting partition version: %s/0", qtid) + } else if found { + // Get the worker(s) responsible for partition 0. + job := partition(qtid.Key(), dax.Partition{ + Num: 0, + Version: v, + }).String() + workers, err := c.TranslateBalancer.WorkersForJobs(ctx, []dax.Job{dax.Job(job)}) + if err != nil { + return errors.Wrapf(err, "getting workers for job: %s", job) + } + + for _, w := range workers { + workerSet.Add(dax.Address(w.ID)) + } + } + + // Get the list of workers responsible for shard data for this table. + if state, err := c.ComputeBalancer.CurrentState(ctx); err != nil { + return errors.Wrap(err, "getting current compute state") + } else { + for _, worker := range state { + for _, job := range worker.Jobs { + if shard, err := decodeShard(job); err != nil { + return errors.Wrapf(err, "decoding shard: %s", job) + } else if shard.table() == qtid.Key() { + workerSet.Add(dax.Address(worker.ID)) + break + } + } + } + } + + // Convert the slice of addresses into a slice of addressMethod containing + // the appropriate method. + addressMethods := applyAddressMethod(workerSet.SortedSlice(), dax.DirectiveMethodDiff) + + // Send a directive to any compute node responsible for this field. + if err := c.sendDirectives(ctx, addressMethods...); err != nil { + return NewErrDirectiveSendFailure(err.Error()) + } + + return nil +} + +func (c *Controller) DropField(ctx context.Context, qtid dax.QualifiedTableID, fldName dax.FieldName) error { + c.mu.Lock() + defer c.mu.Unlock() + + // workerSet maintains the set of workers which have a job assignment change + // and therefore need to be sent an updated Directive. + workerSet := NewAddressSet() + + // If the field has string keys, remove it from the local versionStore. + // TODO: implement RemoveField() on VersionStore interface. + + // Get the worker responsible for partition 0, which handles field key + // translation. Be sure to get the current version. + if v, found, err := c.versionStore.PartitionVersion(ctx, qtid, 0); err != nil { + return errors.Wrapf(err, "getting partition version: %s/0", qtid) + } else if found { + // Get the worker(s) responsible for partition 0. + job := partition(qtid.Key(), dax.Partition{ + Num: 0, + Version: v, + }).String() + workers, err := c.TranslateBalancer.WorkersForJobs(ctx, []dax.Job{dax.Job(job)}) + if err != nil { + return errors.Wrapf(err, "getting workers for job: %s", job) + } + + for _, w := range workers { + workerSet.Add(dax.Address(w.ID)) + } + } + + // Get the list of workers responsible for shard data for this table. + if state, err := c.ComputeBalancer.CurrentState(ctx); err != nil { + return errors.Wrap(err, "getting current compute state") + } else { + for _, worker := range state { + for _, job := range worker.Jobs { + if shard, err := decodeShard(job); err != nil { + return errors.Wrapf(err, "decoding shard: %s", job) + } else if shard.table() == qtid.Key() { + workerSet.Add(dax.Address(worker.ID)) + break + } + } + } + } + + // Convert the slice of addresses into a slice of addressMethod containing + // the appropriate method. + addressMethods := applyAddressMethod(workerSet.SortedSlice(), dax.DirectiveMethodDiff) + + // Send a directive to any compute node responsible for this field. + if err := c.sendDirectives(ctx, addressMethods...); err != nil { + return NewErrDirectiveSendFailure(err.Error()) + } + + return nil +} + +////////////////////////////////// + +func (c *Controller) AddAddresses(ctx context.Context, addrs ...dax.Address) error { + return nil +} + +func (c *Controller) RemoveAddresses(ctx context.Context, addrs ...dax.Address) error { + err := c.DeregisterNodes(ctx, addrs...) + return errors.Wrapf(err, "deregistering nodes: %s", addrs) +} + +func (c *Controller) DebugNodes(ctx context.Context) ([]*dax.Node, error) { + return c.nodeService.Nodes(ctx) +} diff --git a/dax/mds/controller/controller_test.go b/dax/mds/controller/controller_test.go new file mode 100644 index 000000000..eea00a241 --- /dev/null +++ b/dax/mds/controller/controller_test.go @@ -0,0 +1,1373 @@ +package controller_test + +import ( + "context" + "fmt" + "sort" + "sync" + "testing" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/mds/controller" + "github.com/molecula/featurebase/v3/dax/mds/controller/naive/boltdb" + daxtest "github.com/molecula/featurebase/v3/dax/test" + testbolt "github.com/molecula/featurebase/v3/dax/test/boltdb" + "github.com/molecula/featurebase/v3/errors" + "github.com/molecula/featurebase/v3/logger" + "github.com/stretchr/testify/assert" +) + +func TestController(t *testing.T) { + ctx := context.Background() + qual := dax.NewTableQualifier("acme", "db1") + + t.Run("RegisterNode", func(t *testing.T) { + director := newTestDirector() + + schemar, cleanup := daxtest.NewSchemar(t) + defer cleanup() + cfg := controller.Config{ + Director: director, + Schemar: schemar, + } + con := controller.New(cfg) + + // Register a node with an invalid role type. + node0 := &dax.Node{ + Address: "10.0.0.1:80", + RoleTypes: []dax.RoleType{ + "invalid-role-type", + }, + } + err := con.RegisterNodes(ctx, node0) + if assert.Error(t, err) { + assert.True(t, errors.Is(err, controller.ErrCodeRoleTypeInvalid)) + } + + // Register a node with no role type. + node1 := &dax.Node{ + Address: "10.0.0.1:81", + RoleTypes: []dax.RoleType{}, + } + err = con.RegisterNodes(ctx, node1) + if assert.Error(t, err) { + assert.True(t, errors.Is(err, controller.ErrCodeRoleTypeInvalid)) + } + }) + + t.Run("ComputeNodes", func(t *testing.T) { + director := newTestDirector() + schemar, cleanup := daxtest.NewSchemar(t) + defer cleanup() + + db := testbolt.MustOpenDB(t) + db.InitializeBuckets(boltdb.NaiveBalancerBuckets...) + defer func() { + testbolt.MustCloseDB(t, db) + testbolt.CleanupDB(t, db.Path()) + }() + + cfg := controller.Config{ + Director: director, + Schemar: schemar, + BoltDB: db, + StorageMethod: "boltdb", + ComputeBalancer: boltdb.NewBalancer("compute", db, logger.StderrLogger), + TranslateBalancer: boltdb.NewBalancer("translate", db, logger.StderrLogger), + } + con := controller.New(cfg) + + var exp []*dax.Directive + + // Register a node. + node0 := &dax.Node{ + Address: "10.0.0.1:80", + RoleTypes: []dax.RoleType{ + dax.RoleTypeCompute, + }, + } + assert.NoError(t, con.RegisterNodes(ctx, node0)) + + exp = []*dax.Directive{ + { + Address: node0.Address, + Method: dax.DirectiveMethodReset, + Tables: []*dax.QualifiedTable{}, + ComputeRoles: []dax.ComputeRole{}, + TranslateRoles: []dax.TranslateRole{}, + Version: 1, + }, + } + assert.Equal(t, exp, director.flush()) + + // Add a non-keyed table. + tbl0 := daxtest.TestQualifiedTableWithID(t, qual, "2", "foo", 0, false) + assert.NoError(t, schemar.CreateTable(ctx, tbl0)) + assert.NoError(t, con.CreateTable(ctx, tbl0)) + + exp = []*dax.Directive{} + assert.Equal(t, exp, director.flush()) + + // Add the same non-keyed table again. + err := con.CreateTable(ctx, tbl0) + if assert.Error(t, err) { + assert.True(t, errors.Is(err, dax.ErrTableIDExists)) + } + + exp = []*dax.Directive{} + assert.Equal(t, exp, director.flush()) + + // Add a shard. + assert.NoError(t, con.AddShards(ctx, tbl0.QualifiedID(), + dax.NewShard(0, 0), + )) + + exp = []*dax.Directive{ + { + Address: node0.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbl0, + }, + ComputeRoles: []dax.ComputeRole{ + { + TableKey: tbl0.Key(), + Shards: dax.Shards{ + dax.NewShard(0, 0), + }, + }, + }, + TranslateRoles: []dax.TranslateRole{}, + Version: 2, + }, + } + assert.Equal(t, exp, director.flush()) + + // Register two more nodes. + node1 := &dax.Node{ + Address: "10.0.0.1:81", + RoleTypes: []dax.RoleType{ + dax.RoleTypeCompute, + }, + } + assert.NoError(t, con.RegisterNodes(ctx, node1)) + + exp = []*dax.Directive{ + { + Address: node1.Address, + Method: dax.DirectiveMethodReset, + Tables: []*dax.QualifiedTable{}, + ComputeRoles: []dax.ComputeRole{}, + TranslateRoles: []dax.TranslateRole{}, + Version: 3, + }, + } + assert.Equal(t, exp, director.flush()) + + node2 := &dax.Node{ + Address: "10.0.0.1:82", + RoleTypes: []dax.RoleType{ + dax.RoleTypeCompute, + }, + } + assert.NoError(t, con.RegisterNodes(ctx, node2)) + + exp = []*dax.Directive{ + { + Address: node2.Address, + Method: dax.DirectiveMethodReset, + Tables: []*dax.QualifiedTable{}, + ComputeRoles: []dax.ComputeRole{}, + TranslateRoles: []dax.TranslateRole{}, + Version: 4, + }, + } + assert.Equal(t, exp, director.flush()) + + // Add more shards. + assert.NoError(t, con.AddShards(ctx, tbl0.QualifiedID(), + dax.NewShard(1, 0), + dax.NewShard(2, 0), + dax.NewShard(3, 0), + dax.NewShard(5, 0), + dax.NewShard(8, 0), + )) + + exp = []*dax.Directive{ + { + Address: node0.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbl0, + }, + ComputeRoles: []dax.ComputeRole{ + { + TableKey: tbl0.Key(), + Shards: dax.Shards{ + dax.NewShard(0, 0), + dax.NewShard(3, 0), + }, + }, + }, + TranslateRoles: []dax.TranslateRole{}, + Version: 5, + }, + { + Address: node1.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbl0, + }, + ComputeRoles: []dax.ComputeRole{ + { + TableKey: tbl0.Key(), + Shards: dax.Shards{ + dax.NewShard(1, 0), + dax.NewShard(5, 0), + }, + }, + }, + TranslateRoles: []dax.TranslateRole{}, + Version: 6, + }, + { + Address: node2.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbl0, + }, + ComputeRoles: []dax.ComputeRole{ + { + TableKey: tbl0.Key(), + Shards: dax.Shards{ + dax.NewShard(2, 0), + dax.NewShard(8, 0), + }, + }, + }, + TranslateRoles: []dax.TranslateRole{}, + Version: 7, + }, + } + assert.Equal(t, exp, director.flush()) + + // Add another non-keyed table. + tbl1 := daxtest.TestQualifiedTableWithID(t, qual, "1", "bar", 0, false) + assert.NoError(t, schemar.CreateTable(ctx, tbl1)) + assert.NoError(t, con.CreateTable(ctx, tbl1)) + + // Add more shards. + assert.NoError(t, con.AddShards(ctx, tbl1.QualifiedID(), + dax.NewShard(3, 0), + dax.NewShard(5, 0), + dax.NewShard(8, 0), + dax.NewShard(13, 0), + )) + + exp = []*dax.Directive{ + { + Address: node0.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbl1, + tbl0, + }, + ComputeRoles: []dax.ComputeRole{ + { + TableKey: tbl1.Key(), + Shards: dax.Shards{ + dax.NewShard(3, 0), + dax.NewShard(13, 0), + }, + }, + { + TableKey: tbl0.Key(), + Shards: dax.Shards{ + dax.NewShard(0, 0), + dax.NewShard(3, 0), + }, + }, + }, + TranslateRoles: []dax.TranslateRole{}, + Version: 8, + }, + { + Address: node1.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbl1, + tbl0, + }, + ComputeRoles: []dax.ComputeRole{ + { + TableKey: tbl1.Key(), + Shards: dax.Shards{ + dax.NewShard(5, 0), + }, + }, + { + TableKey: tbl0.Key(), + Shards: dax.Shards{ + dax.NewShard(1, 0), + dax.NewShard(5, 0), + }, + }, + }, + TranslateRoles: []dax.TranslateRole{}, + Version: 9, + }, + { + Address: node2.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbl1, + tbl0, + }, + ComputeRoles: []dax.ComputeRole{ + { + TableKey: tbl1.Key(), + Shards: dax.Shards{ + dax.NewShard(8, 0), + }, + }, + { + TableKey: tbl0.Key(), + Shards: dax.Shards{ + dax.NewShard(2, 0), + dax.NewShard(8, 0), + }, + }, + }, + TranslateRoles: []dax.TranslateRole{}, + Version: 10, + }, + } + assert.Equal(t, exp, director.flush()) + + // Remove a node. + assert.NoError(t, con.DeregisterNodes(ctx, node1.Address)) + + exp = []*dax.Directive{ + { + Address: node0.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbl1, + tbl0, + }, + ComputeRoles: []dax.ComputeRole{ + { + TableKey: tbl1.Key(), + Shards: dax.Shards{ + dax.NewShard(3, 0), + dax.NewShard(13, 0), + }, + }, + { + TableKey: tbl0.Key(), + Shards: dax.Shards{ + dax.NewShard(0, 0), + dax.NewShard(1, 0), + dax.NewShard(3, 0), + }, + }, + }, + TranslateRoles: []dax.TranslateRole{}, + Version: 11, + }, + { + Address: node2.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbl1, + tbl0, + }, + ComputeRoles: []dax.ComputeRole{ + { + TableKey: tbl1.Key(), + Shards: dax.Shards{ + dax.NewShard(5, 0), + dax.NewShard(8, 0), + }, + }, + { + TableKey: tbl0.Key(), + Shards: dax.Shards{ + dax.NewShard(2, 0), + dax.NewShard(5, 0), + dax.NewShard(8, 0), + }, + }, + }, + TranslateRoles: []dax.TranslateRole{}, + Version: 12, + }, + } + assert.Equal(t, exp, director.flush()) + + // Remove another node. + assert.NoError(t, con.DeregisterNodes(ctx, node0.Address)) + + exp = []*dax.Directive{ + { + Address: node2.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbl1, + tbl0, + }, + ComputeRoles: []dax.ComputeRole{ + { + TableKey: tbl1.Key(), + Shards: dax.Shards{ + dax.NewShard(3, 0), + dax.NewShard(5, 0), + dax.NewShard(8, 0), + dax.NewShard(13, 0), + }, + }, + { + TableKey: tbl0.Key(), + Shards: dax.Shards{ + dax.NewShard(0, 0), + dax.NewShard(1, 0), + dax.NewShard(2, 0), + dax.NewShard(3, 0), + dax.NewShard(5, 0), + dax.NewShard(8, 0), + }, + }, + }, + TranslateRoles: []dax.TranslateRole{}, + Version: 13, + }, + } + assert.Equal(t, exp, director.flush()) + + // Remove final node. + assert.NoError(t, con.DeregisterNodes(ctx, node2.Address)) + + exp = []*dax.Directive{} + assert.Equal(t, exp, director.flush()) + + // Add a new node and ensure that the free shards get assigned to it. + node3 := &dax.Node{ + Address: "10.0.0.1:83", + RoleTypes: []dax.RoleType{ + dax.RoleTypeCompute, + }, + } + assert.NoError(t, con.RegisterNodes(ctx, node3)) + + exp = []*dax.Directive{ + { + Address: node3.Address, + Method: dax.DirectiveMethodReset, + Tables: []*dax.QualifiedTable{ + tbl1, + tbl0, + }, + ComputeRoles: []dax.ComputeRole{ + { + TableKey: tbl1.Key(), + Shards: dax.Shards{ + dax.NewShard(3, 0), + dax.NewShard(5, 0), + dax.NewShard(8, 0), + dax.NewShard(13, 0), + }, + }, + { + TableKey: tbl0.Key(), + Shards: dax.Shards{ + dax.NewShard(0, 0), + dax.NewShard(1, 0), + dax.NewShard(2, 0), + dax.NewShard(3, 0), + dax.NewShard(5, 0), + dax.NewShard(8, 0), + }, + }, + }, + TranslateRoles: []dax.TranslateRole{}, + Version: 14, + }, + } + assert.Equal(t, exp, director.flush()) + + // Remove shards. + assert.NoError(t, con.RemoveShards(ctx, tbl0.QualifiedID(), + dax.NewShard(2, 0), + dax.NewShard(5, 0), + )) + + exp = []*dax.Directive{ + { + Address: node3.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbl1, + tbl0, + }, + ComputeRoles: []dax.ComputeRole{ + { + TableKey: tbl1.Key(), + Shards: dax.Shards{ + dax.NewShard(3, 0), + dax.NewShard(5, 0), + dax.NewShard(8, 0), + dax.NewShard(13, 0), + }, + }, + { + TableKey: tbl0.Key(), + Shards: dax.Shards{ + dax.NewShard(0, 0), + dax.NewShard(1, 0), + dax.NewShard(3, 0), + dax.NewShard(8, 0), + }, + }, + }, + TranslateRoles: []dax.TranslateRole{}, + Version: 15, + }, + } + assert.Equal(t, exp, director.flush()) + + // Remove shards, one which does not exist. + // Currently that doesn't result in an error, it simply no-ops on trying + // to remove 99. + assert.NoError(t, con.RemoveShards(ctx, tbl0.QualifiedID(), + dax.NewShard(3, 0), + dax.NewShard(99, 0), + )) + + exp = []*dax.Directive{ + { + Address: node3.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbl1, + tbl0, + }, + ComputeRoles: []dax.ComputeRole{ + { + TableKey: tbl1.Key(), + Shards: dax.Shards{ + dax.NewShard(3, 0), + dax.NewShard(5, 0), + dax.NewShard(8, 0), + dax.NewShard(13, 0), + }, + }, + { + TableKey: tbl0.Key(), + Shards: dax.Shards{ + dax.NewShard(0, 0), + dax.NewShard(1, 0), + dax.NewShard(8, 0), + }, + }, + }, + TranslateRoles: []dax.TranslateRole{}, + Version: 16, + }, + } + assert.Equal(t, exp, director.flush()) + + // Remove a table. + assert.NoError(t, con.DropTable(ctx, tbl0.QualifiedID())) + + exp = []*dax.Directive{ + { + Address: node3.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbl1, + }, + ComputeRoles: []dax.ComputeRole{ + { + TableKey: tbl1.Key(), + Shards: dax.Shards{ + dax.NewShard(3, 0), + dax.NewShard(5, 0), + dax.NewShard(8, 0), + dax.NewShard(13, 0), + }, + }, + }, + TranslateRoles: []dax.TranslateRole{}, + Version: 17, + }, + } + assert.Equal(t, exp, director.flush()) + + // Remove a node which doesn't exist. + err = con.DeregisterNodes(ctx, "invalidNode") + if assert.Error(t, err) { + assert.True(t, errors.Is(err, dax.ErrNodeDoesNotExist)) + } + + }) + + t.Run("TranslateNodes", func(t *testing.T) { + invalidQtid := dax.NewQualifiedTableID( + dax.NewTableQualifier("", ""), + dax.TableID("invalidID"), + ) + + director := newTestDirector() + schemar, cleanup := daxtest.NewSchemar(t) + defer cleanup() + + db := testbolt.MustOpenDB(t) + db.InitializeBuckets(boltdb.NaiveBalancerBuckets...) + defer func() { + testbolt.MustCloseDB(t, db) + testbolt.CleanupDB(t, db.Path()) + }() + + cfg := controller.Config{ + Director: director, + Schemar: schemar, + BoltDB: db, + StorageMethod: "boltdb", + ComputeBalancer: boltdb.NewBalancer("compute", db, logger.StderrLogger), + TranslateBalancer: boltdb.NewBalancer("translate", db, logger.StderrLogger), + } + con := controller.New(cfg) + + var exp []*dax.Directive + + // Register a node. + node0 := &dax.Node{ + Address: "10.0.0.1:80", + RoleTypes: []dax.RoleType{ + dax.RoleTypeTranslate, + }, + } + assert.NoError(t, con.RegisterNodes(ctx, node0)) + + exp = []*dax.Directive{ + { + Address: node0.Address, + Method: dax.DirectiveMethodReset, + Tables: []*dax.QualifiedTable{}, + ComputeRoles: []dax.ComputeRole{}, + TranslateRoles: []dax.TranslateRole{}, + Version: 1, + }, + } + assert.Equal(t, exp, director.flush()) + + // Try registering the same node. This should be ok. + assert.NoError(t, con.RegisterNodes(ctx, node0)) + + exp = []*dax.Directive{} + assert.Equal(t, exp, director.flush()) + + // Add a keyed table. + tbl0 := daxtest.TestQualifiedTableWithID(t, qual, "2", "foo", 8, true) + assert.NoError(t, schemar.CreateTable(ctx, tbl0)) + assert.NoError(t, con.CreateTable(ctx, tbl0)) + + // Check directives. + exp = []*dax.Directive{ + { + Address: node0.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbl0, + }, + ComputeRoles: []dax.ComputeRole{}, + TranslateRoles: []dax.TranslateRole{ + { + TableKey: tbl0.Key(), + Partitions: dax.Partitions{ + dax.NewPartition(0, 0), + dax.NewPartition(1, 0), + dax.NewPartition(2, 0), + dax.NewPartition(3, 0), + dax.NewPartition(4, 0), + dax.NewPartition(5, 0), + dax.NewPartition(6, 0), + dax.NewPartition(7, 0), + }, + }, + }, + Version: 2, + }, + } + assert.Equal(t, exp, director.flush()) + + // Register two more nodes. + node1 := &dax.Node{ + Address: "10.0.0.1:81", + RoleTypes: []dax.RoleType{ + dax.RoleTypeTranslate, + }, + } + assert.NoError(t, con.RegisterNodes(ctx, node1)) + + exp = []*dax.Directive{ + { + Address: node0.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbl0, + }, + ComputeRoles: []dax.ComputeRole{}, + TranslateRoles: []dax.TranslateRole{ + { + TableKey: tbl0.Key(), + Partitions: dax.Partitions{ + dax.NewPartition(0, 0), + dax.NewPartition(1, 0), + dax.NewPartition(2, 0), + dax.NewPartition(3, 0), + }, + }, + }, + Version: 3, + }, + { + Address: node1.Address, + Method: dax.DirectiveMethodReset, + Tables: []*dax.QualifiedTable{ + tbl0, + }, + ComputeRoles: []dax.ComputeRole{}, + TranslateRoles: []dax.TranslateRole{ + { + TableKey: tbl0.Key(), + Partitions: dax.Partitions{ + dax.NewPartition(4, 0), + dax.NewPartition(5, 0), + dax.NewPartition(6, 0), + dax.NewPartition(7, 0), + }, + }, + }, + Version: 4, + }, + } + assert.Equal(t, exp, director.flush()) + + node2 := &dax.Node{ + Address: "10.0.0.1:82", + RoleTypes: []dax.RoleType{ + dax.RoleTypeTranslate, + }, + } + assert.NoError(t, con.RegisterNodes(ctx, node2)) + + exp = []*dax.Directive{ + { + Address: node0.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbl0, + }, + ComputeRoles: []dax.ComputeRole{}, + TranslateRoles: []dax.TranslateRole{ + { + TableKey: tbl0.Key(), + Partitions: dax.Partitions{ + dax.NewPartition(0, 0), + dax.NewPartition(1, 0), + dax.NewPartition(2, 0), + }, + }, + }, + Version: 5, + }, + { + Address: node1.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbl0, + }, + ComputeRoles: []dax.ComputeRole{}, + TranslateRoles: []dax.TranslateRole{ + { + TableKey: tbl0.Key(), + Partitions: dax.Partitions{ + dax.NewPartition(4, 0), + dax.NewPartition(5, 0), + dax.NewPartition(6, 0), + }, + }, + }, + Version: 6, + }, + { + Address: node2.Address, + Method: dax.DirectiveMethodReset, + Tables: []*dax.QualifiedTable{ + tbl0, + }, + ComputeRoles: []dax.ComputeRole{}, + TranslateRoles: []dax.TranslateRole{ + { + TableKey: tbl0.Key(), + Partitions: dax.Partitions{ + dax.NewPartition(3, 0), + dax.NewPartition(7, 0), + }, + }, + }, + Version: 7, + }, + } + assert.Equal(t, exp, director.flush()) + + // Add another keyed table. + // Make PartitionN double digit to ensure that partition ints aren't + // sorted as strings. Also, it should be large enough to spill over + // onto node0. + tbl1 := daxtest.TestQualifiedTableWithID(t, qual, "1", "bar", 24, true) + assert.NoError(t, schemar.CreateTable(ctx, tbl1)) + assert.NoError(t, con.CreateTable(ctx, tbl1)) + + // Check directives. + exp = []*dax.Directive{ + { + Address: node0.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbl1, + tbl0, + }, + ComputeRoles: []dax.ComputeRole{}, + TranslateRoles: []dax.TranslateRole{ + { + TableKey: tbl1.Key(), + Partitions: dax.Partitions{ + dax.NewPartition(1, 0), + dax.NewPartition(4, 0), + dax.NewPartition(7, 0), + dax.NewPartition(10, 0), + dax.NewPartition(13, 0), + dax.NewPartition(16, 0), + dax.NewPartition(19, 0), + dax.NewPartition(22, 0), + }, + }, + { + TableKey: tbl0.Key(), + Partitions: dax.Partitions{ + dax.NewPartition(0, 0), + dax.NewPartition(1, 0), + dax.NewPartition(2, 0), + }, + }, + }, + Version: 8, + }, + { + Address: node1.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbl1, + tbl0, + }, + ComputeRoles: []dax.ComputeRole{}, + TranslateRoles: []dax.TranslateRole{ + { + TableKey: tbl1.Key(), + Partitions: dax.Partitions{ + dax.NewPartition(2, 0), + dax.NewPartition(5, 0), + dax.NewPartition(8, 0), + dax.NewPartition(11, 0), + dax.NewPartition(14, 0), + dax.NewPartition(17, 0), + dax.NewPartition(20, 0), + dax.NewPartition(23, 0), + }, + }, + { + TableKey: tbl0.Key(), + Partitions: dax.Partitions{ + dax.NewPartition(4, 0), + dax.NewPartition(5, 0), + dax.NewPartition(6, 0), + }, + }, + }, + Version: 9, + }, + { + Address: node2.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbl1, + tbl0, + }, + ComputeRoles: []dax.ComputeRole{}, + TranslateRoles: []dax.TranslateRole{ + { + TableKey: tbl1.Key(), + Partitions: dax.Partitions{ + dax.NewPartition(0, 0), + dax.NewPartition(3, 0), + dax.NewPartition(6, 0), + dax.NewPartition(9, 0), + dax.NewPartition(12, 0), + dax.NewPartition(15, 0), + dax.NewPartition(18, 0), + dax.NewPartition(21, 0), + }, + }, + { + TableKey: tbl0.Key(), + Partitions: dax.Partitions{ + dax.NewPartition(3, 0), + dax.NewPartition(7, 0), + }, + }, + }, + Version: 10, + }, + } + assert.Equal(t, exp, director.flush()) + + // Remove a keyed table. + assert.NoError(t, con.DropTable(ctx, tbl0.QualifiedID())) + + // Check directives. + exp = []*dax.Directive{ + { + Address: node0.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbl1, + }, + ComputeRoles: []dax.ComputeRole{}, + TranslateRoles: []dax.TranslateRole{ + { + TableKey: tbl1.Key(), + Partitions: dax.Partitions{ + dax.NewPartition(1, 0), + dax.NewPartition(4, 0), + dax.NewPartition(7, 0), + dax.NewPartition(10, 0), + dax.NewPartition(13, 0), + dax.NewPartition(16, 0), + dax.NewPartition(19, 0), + dax.NewPartition(22, 0), + }, + }, + }, + Version: 11, + }, + { + Address: node1.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbl1, + }, + ComputeRoles: []dax.ComputeRole{}, + TranslateRoles: []dax.TranslateRole{ + { + TableKey: tbl1.Key(), + Partitions: dax.Partitions{ + dax.NewPartition(2, 0), + dax.NewPartition(5, 0), + dax.NewPartition(8, 0), + dax.NewPartition(11, 0), + dax.NewPartition(14, 0), + dax.NewPartition(17, 0), + dax.NewPartition(20, 0), + dax.NewPartition(23, 0), + }, + }, + }, + Version: 12, + }, + { + Address: node2.Address, + Method: dax.DirectiveMethodDiff, + Tables: []*dax.QualifiedTable{ + tbl1, + }, + ComputeRoles: []dax.ComputeRole{}, + TranslateRoles: []dax.TranslateRole{ + { + TableKey: tbl1.Key(), + Partitions: dax.Partitions{ + dax.NewPartition(0, 0), + dax.NewPartition(3, 0), + dax.NewPartition(6, 0), + dax.NewPartition(9, 0), + dax.NewPartition(12, 0), + dax.NewPartition(15, 0), + dax.NewPartition(18, 0), + dax.NewPartition(21, 0), + }, + }, + }, + Version: 13, + }, + } + assert.Equal(t, exp, director.flush()) + + // Remove a table which doesn't exist. + err := con.DropTable(ctx, invalidQtid) + if assert.Error(t, err) { + assert.True(t, errors.Is(err, dax.ErrTableIDDoesNotExist)) + } + + // Add shards to a table which doesn't exist. + err = con.AddShards(ctx, invalidQtid, + dax.NewShard(1, 0), + dax.NewShard(2, 0), + ) + if assert.Error(t, err) { + assert.True(t, errors.Is(err, dax.ErrTableIDDoesNotExist)) + } + + // Register an invalid node. + nodeX := &dax.Node{ + Address: "", + } + err = con.RegisterNodes(ctx, nodeX) + if assert.Error(t, err) { + assert.True(t, errors.Is(err, controller.ErrCodeNodeKeyInvalid)) + } + }) + + t.Run("GetNodes", func(t *testing.T) { + schemar, cleanup := daxtest.NewSchemar(t) + defer cleanup() + + db := testbolt.MustOpenDB(t) + db.InitializeBuckets(boltdb.NaiveBalancerBuckets...) + defer func() { + testbolt.MustCloseDB(t, db) + testbolt.CleanupDB(t, db.Path()) + }() + + cfg := controller.Config{ + Schemar: schemar, + BoltDB: db, + StorageMethod: "boltdb", + ComputeBalancer: boltdb.NewBalancer("compute", db, logger.StderrLogger), + TranslateBalancer: boltdb.NewBalancer("translate", db, logger.StderrLogger), + } + con := controller.New(cfg) + + // Register two nodes. + node0 := &dax.Node{ + Address: "10.0.0.1:80", + RoleTypes: []dax.RoleType{ + dax.RoleTypeCompute, + dax.RoleTypeTranslate, + }, + } + assert.NoError(t, con.RegisterNodes(ctx, node0)) + node1 := &dax.Node{ + Address: "10.0.0.1:81", + RoleTypes: []dax.RoleType{ + dax.RoleTypeCompute, + dax.RoleTypeTranslate, + }, + } + assert.NoError(t, con.RegisterNodes(ctx, node1)) + + // Add a keyed table. + tbl0 := daxtest.TestQualifiedTable(t, qual, "foo", 12, true) + assert.NoError(t, schemar.CreateTable(ctx, tbl0)) + assert.NoError(t, con.CreateTable(ctx, tbl0)) + + // Add shards. + assert.NoError(t, con.AddShards(ctx, tbl0.QualifiedID(), + dax.NewShard(0, 0), + dax.NewShard(1, 0), + dax.NewShard(2, 0), + dax.NewShard(3, 0), + dax.NewShard(11, 0), + dax.NewShard(12, 0), + )) + + t.Run("ComputeRole", func(t *testing.T) { + tests := []struct { + role dax.Role + isWrite bool + exp []dax.AssignedNode + }{ + { + role: &dax.ComputeRole{ + TableKey: tbl0.Key(), + Shards: dax.NewShards(0, 1, 2, 3), + }, + exp: []dax.AssignedNode{ + { + Address: node0.Address, + Role: &dax.ComputeRole{ + TableKey: tbl0.Key(), + Shards: dax.Shards{ + dax.NewShard(0, 0), + dax.NewShard(2, 0), + }, + }, + }, + { + Address: node1.Address, + Role: &dax.ComputeRole{ + TableKey: tbl0.Key(), + Shards: dax.Shards{ + dax.NewShard(1, 0), + dax.NewShard(3, 0), + }, + }, + }, + }, + }, + { + role: &dax.ComputeRole{ + TableKey: tbl0.Key(), + Shards: dax.NewShards(1), + }, + exp: []dax.AssignedNode{ + { + Address: node1.Address, + Role: &dax.ComputeRole{ + TableKey: tbl0.Key(), + Shards: dax.Shards{ + dax.NewShard(1, 0), + }, + }, + }, + }, + }, + { + // Add unassigned shards. + role: &dax.ComputeRole{ + TableKey: tbl0.Key(), + Shards: dax.NewShards(1, 888, 889), + }, + isWrite: true, + exp: []dax.AssignedNode{ + { + Address: node0.Address, + Role: &dax.ComputeRole{ + TableKey: tbl0.Key(), + Shards: dax.Shards{ + dax.NewShard(888, 0), + }, + }, + }, + { + Address: node1.Address, + Role: &dax.ComputeRole{ + TableKey: tbl0.Key(), + Shards: dax.Shards{ + dax.NewShard(1, 0), + dax.NewShard(889, 0), + }, + }, + }, + }, + }, + { + // Ensure shards are not returned sorted as strings. + role: &dax.ComputeRole{ + TableKey: tbl0.Key(), + Shards: dax.NewShards(2, 11), + }, + exp: []dax.AssignedNode{ + { + Address: node0.Address, + Role: &dax.ComputeRole{ + TableKey: tbl0.Key(), + Shards: dax.Shards{ + dax.NewShard(2, 0), + dax.NewShard(11, 0), + }, + }, + }, + }, + }, + } + for i, test := range tests { + t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) { + nodes, err := con.Nodes(ctx, test.role, test.isWrite) + assert.NoError(t, err) + assert.Equal(t, test.exp, nodes) + }) + } + }) + + t.Run("TranslateRole", func(t *testing.T) { + tests := []struct { + role dax.Role + isWrite bool + exp []dax.AssignedNode + expErrCode errors.Code + }{ + { + role: &dax.TranslateRole{ + TableKey: tbl0.Key(), + Partitions: dax.Partitions{ + dax.NewPartition(0, -1), + }, + }, + isWrite: true, + exp: []dax.AssignedNode{ + { + Address: node0.Address, + Role: &dax.TranslateRole{ + TableKey: tbl0.Key(), + Partitions: dax.Partitions{ + dax.NewPartition(0, 0), + }, + }, + }, + }, + }, + { + role: &dax.TranslateRole{ + TableKey: tbl0.Key(), + Partitions: dax.Partitions{ + dax.NewPartition(0, -1), + dax.NewPartition(1, -1), + dax.NewPartition(2, -1), + dax.NewPartition(3, -1), + dax.NewPartition(999, -1), + }, + }, + isWrite: false, + exp: []dax.AssignedNode{ + { + Address: node0.Address, + Role: &dax.TranslateRole{ + TableKey: tbl0.Key(), + Partitions: dax.Partitions{ + dax.NewPartition(0, 0), + dax.NewPartition(2, 0), + }, + }, + }, + { + Address: node1.Address, + Role: &dax.TranslateRole{ + TableKey: tbl0.Key(), + Partitions: dax.Partitions{ + dax.NewPartition(1, 0), + dax.NewPartition(3, 0), + }, + }, + }, + }, + }, + { + role: &dax.TranslateRole{ + TableKey: tbl0.Key(), + Partitions: dax.Partitions{ + dax.NewPartition(1, -1), + }, + }, + isWrite: false, + exp: []dax.AssignedNode{ + { + Address: node1.Address, + Role: &dax.TranslateRole{ + TableKey: tbl0.Key(), + Partitions: dax.Partitions{ + dax.NewPartition(1, 0), + }, + }, + }, + }, + }, + { + // Ensure partitions are not returned sorted as strings. + role: &dax.TranslateRole{ + TableKey: tbl0.Key(), + Partitions: dax.Partitions{ + dax.NewPartition(2, -1), + dax.NewPartition(10, -1), + }, + }, + isWrite: false, + exp: []dax.AssignedNode{ + { + Address: node0.Address, + Role: &dax.TranslateRole{ + TableKey: tbl0.Key(), + Partitions: dax.Partitions{ + dax.NewPartition(2, 0), + dax.NewPartition(10, 0), + }, + }, + }, + }, + }, + } + for i, test := range tests { + t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) { + nodes, err := con.Nodes(ctx, test.role, test.isWrite) + + if test.expErrCode != "" { + assert.True(t, errors.Is(err, test.expErrCode)) + } else { + assert.NoError(t, err) + assert.Equal(t, test.exp, nodes) + } + }) + } + }) + }) +} + +////////////////////////////////////////////////////// + +// Ensure type implements interface. +var _ controller.Director = &testDirector{} + +// testDirector is an implementation of the Director interface used for testing. +type testDirector struct { + mu sync.Mutex + dirs []*dax.Directive +} + +func newTestDirector() *testDirector { + return &testDirector{} +} + +func (d *testDirector) SendDirective(ctx context.Context, dir *dax.Directive) error { + d.mu.Lock() + defer d.mu.Unlock() + d.dirs = append(d.dirs, dir) + return nil +} + +func (d *testDirector) SendSnapshotShardDataRequest(ctx context.Context, req *dax.SnapshotShardDataRequest) error { + return nil +} + +func (d *testDirector) SendSnapshotTableKeysRequest(ctx context.Context, req *dax.SnapshotTableKeysRequest) error { + return nil +} + +func (d *testDirector) SendSnapshotFieldKeysRequest(ctx context.Context, req *dax.SnapshotFieldKeysRequest) error { + return nil +} + +// flush returns all the directives that have been captured through the Send() +// method and then resets the internal list. +func (d *testDirector) flush() []*dax.Directive { + out := make([]*dax.Directive, len(d.dirs)) + copy(out, d.dirs) + + // Zero out the slice (but retain allocated memory). + d.dirs = d.dirs[:0] + + // Since the directives can be received asyncronously, sort them here so + // that we can more easily compare them in tests. + sort.Sort(dax.Directives(out)) + + return out +} diff --git a/dax/mds/controller/director.go b/dax/mds/controller/director.go new file mode 100644 index 000000000..29710e137 --- /dev/null +++ b/dax/mds/controller/director.go @@ -0,0 +1,40 @@ +package controller + +import ( + "context" + + "github.com/molecula/featurebase/v3/dax" +) + +type Director interface { + SendDirective(ctx context.Context, dir *dax.Directive) error + SendSnapshotShardDataRequest(ctx context.Context, req *dax.SnapshotShardDataRequest) error + SendSnapshotTableKeysRequest(ctx context.Context, req *dax.SnapshotTableKeysRequest) error + SendSnapshotFieldKeysRequest(ctx context.Context, req *dax.SnapshotFieldKeysRequest) error +} + +// Ensure type implements interface. +var _ Director = &NopDirector{} + +// NopDirector is a no-op implementation of the Director interface. +type NopDirector struct{} + +func NewNopDirector() *NopDirector { + return &NopDirector{} +} + +func (d *NopDirector) SendDirective(ctx context.Context, dir *dax.Directive) error { + return nil +} + +func (d *NopDirector) SendSnapshotShardDataRequest(ctx context.Context, req *dax.SnapshotShardDataRequest) error { + return nil +} + +func (d *NopDirector) SendSnapshotTableKeysRequest(ctx context.Context, req *dax.SnapshotTableKeysRequest) error { + return nil +} + +func (d *NopDirector) SendSnapshotFieldKeysRequest(ctx context.Context, req *dax.SnapshotFieldKeysRequest) error { + return nil +} diff --git a/dax/mds/controller/errors.go b/dax/mds/controller/errors.go new file mode 100644 index 000000000..f4c13219b --- /dev/null +++ b/dax/mds/controller/errors.go @@ -0,0 +1,101 @@ +package controller + +import ( + "fmt" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/errors" +) + +const ( + // ErrCodeCustom can be used to return a custom error message. + ErrCodeCustom errors.Code = "CustomError" + + // ErrCodeInternal can be used when the cause of an error can't be + // determined. It can be accompanied by a single string message. + ErrCodeInternal errors.Code = "InternalError" + + // ErrCodeTODO can be used as a placeholder until a proper error code is + // created and assigned. + ErrCodeTODO errors.Code = "TODOError" + + ErrCodeNodeExists errors.Code = "NodeExists" + ErrCodeNodeKeyInvalid errors.Code = "NodeKeyInvalid" + ErrCodeNoAvailableNode errors.Code = "NoAvailableNode" + + ErrCodeRoleTypeInvalid errors.Code = "RoleTypeInvalid" + + ErrCodeDirectiveSendFailure errors.Code = "DirectiveSendFailure" + + ErrCodeInvalidRequest errors.Code = "InvalidRequest" + + ErrCodeUnassignedJobs errors.Code = "UnassignedJobs" + + UndefinedErrorMessage string = "undefined message format" +) + +// NewErrCustom can be used to return a custom error message. +func NewErrCustom() error { + return errors.New( + ErrCodeCustom, + "", + ) +} + +// NewErrInternal can be used when the cause of an error can't be determined. It +// can be accompanied by a single string message. +func NewErrInternal(msg string) error { + return errors.New( + ErrCodeInternal, + fmt.Sprintf("internal error: %s", msg), + ) +} + +func NewErrNodeExists(addr dax.Address) error { + return errors.New( + ErrCodeNodeExists, + fmt.Sprintf("node '%s' already exists", addr), + ) +} + +func NewErrNodeKeyInvalid(addr dax.Address) error { + return errors.New( + ErrCodeNodeKeyInvalid, + fmt.Sprintf("node key '%s' is invalid", addr), + ) +} + +func NewErrNoAvailableNode() error { + return errors.New( + ErrCodeNoAvailableNode, + "no available node", + ) +} + +func NewErrRoleTypeInvalid(roleType dax.RoleType) error { + return errors.New( + ErrCodeRoleTypeInvalid, + fmt.Sprintf("role type '%s' is invalid", roleType), + ) +} + +func NewErrDirectiveSendFailure(msg string) error { + return errors.New( + ErrCodeDirectiveSendFailure, + fmt.Sprintf("directive failed to send: %s", msg), + ) +} + +func NewErrInvalidRequest(msg string) error { + return errors.New( + ErrCodeInvalidRequest, + fmt.Sprintf("invalid request: %s", msg), + ) +} + +func NewErrUnassignedJobs(jobs []dax.Job) error { + return errors.New( + ErrCodeUnassignedJobs, + fmt.Sprintf("found %d unassigned jobs", len(jobs)), + ) +} diff --git a/dax/mds/controller/http/director.go b/dax/mds/controller/http/director.go new file mode 100644 index 000000000..2560aa440 --- /dev/null +++ b/dax/mds/controller/http/director.go @@ -0,0 +1,185 @@ +// Package http provides the http implementation of the Director interface. +package http + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net" + "net/http" + "time" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/errors" + "github.com/molecula/featurebase/v3/logger" +) + +// Director is an http implementation of the Director interface. +type Director struct { + // directivePath is the path portion of the URI to which directives should + // be POSTed. + directivePath string + + // snapshotRequestPath is the path portion of the URI to which snapshot + // requests should be POSTed. + snapshotRequestPath string + + client *http.Client + + logger logger.Logger +} + +func NewDirector(cfg DirectorConfig) *Director { + var logr logger.Logger = logger.NopLogger + if cfg.Logger != nil { + logr = cfg.Logger + } + + return &Director{ + directivePath: cfg.DirectivePath, + snapshotRequestPath: cfg.SnapshotRequestPath, + logger: logr, + client: &http.Client{ + Transport: &http.Transport{ + Proxy: http.ProxyFromEnvironment, + DialContext: (&net.Dialer{ + Timeout: 2 * time.Second, + KeepAlive: 30 * time.Second, + }).DialContext, + ForceAttemptHTTP2: true, + MaxIdleConns: 100, + IdleConnTimeout: 90 * time.Second, + TLSHandshakeTimeout: 3 * time.Second, + ExpectContinueTimeout: 1 * time.Second, + }, + }, + } +} + +type DirectorConfig struct { + DirectivePath string + SnapshotRequestPath string + Logger logger.Logger +} + +func (d *Director) SendDirective(ctx context.Context, dir *dax.Directive) error { + url := fmt.Sprintf("%s/%s", dir.Address.WithScheme("http"), d.directivePath) + d.logger.Printf("SEND HTTP directive to: %s\n", url) + + // Encode the request. + postBody, err := json.Marshal(dir) + if err != nil { + return errors.Wrap(err, "marshalling directive to json") + } + requestBody := bytes.NewBuffer(postBody) + + // Post the request. + request, _ := http.NewRequest(http.MethodPost, url, requestBody) + request.Header.Add("Content-Type", "application/json") + request.Header.Add("Accept", "application/json") + + resp, err := d.client.Do(request) + if err != nil { + return errors.Wrap(err, "doing send directive") + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(resp.Body) + return errors.Errorf("status code: %d: %s", resp.StatusCode, b) + } + + return nil +} + +func (d *Director) SendSnapshotShardDataRequest(ctx context.Context, req *dax.SnapshotShardDataRequest) error { + url := fmt.Sprintf("%s/%s/shard-data", req.Address.WithScheme("http"), d.snapshotRequestPath) + d.logger.Printf("SEND HTTP snapshot shard data request to: %s\n", url) + + // Encode the request. + postBody, err := json.Marshal(req) + if err != nil { + return errors.Wrap(err, "marshalling snapshot shard data request to json") + } + requestBody := bytes.NewBuffer(postBody) + + // Post the request. + request, _ := http.NewRequest(http.MethodPost, url, requestBody) + request.Header.Add("Content-Type", "application/json") + request.Header.Add("Accept", "application/json") + + resp, err := d.client.Do(request) + if err != nil { + return errors.Wrap(err, "doing snapshot shard data request") + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(resp.Body) + return errors.Errorf("status code: %d: %s", resp.StatusCode, b) + } + + return nil +} + +func (d *Director) SendSnapshotTableKeysRequest(ctx context.Context, req *dax.SnapshotTableKeysRequest) error { + url := fmt.Sprintf("%s/%s/table-keys", req.Address.WithScheme("http"), d.snapshotRequestPath) + d.logger.Printf("SEND HTTP snapshot table keys request to: %s\n", url) + + // Encode the request. + postBody, err := json.Marshal(req) + if err != nil { + return errors.Wrap(err, "marshalling snapshot table keys request to json") + } + requestBody := bytes.NewBuffer(postBody) + + // Post the request. + request, _ := http.NewRequest(http.MethodPost, url, requestBody) + request.Header.Add("Content-Type", "application/json") + request.Header.Add("Accept", "application/json") + + resp, err := d.client.Do(request) + if err != nil { + return errors.Wrap(err, "doing snapshot table keys request") + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(resp.Body) + return errors.Errorf("status code: %d: %s", resp.StatusCode, b) + } + + return nil +} + +func (d *Director) SendSnapshotFieldKeysRequest(ctx context.Context, req *dax.SnapshotFieldKeysRequest) error { + url := fmt.Sprintf("%s/%s/field-keys", req.Address.WithScheme("http"), d.snapshotRequestPath) + d.logger.Printf("SEND HTTP snapshot field keys request to: %s\n", url) + + // Encode the request. + postBody, err := json.Marshal(req) + if err != nil { + return errors.Wrap(err, "marshalling snapshot field keys request to json") + } + requestBody := bytes.NewBuffer(postBody) + + // Post the request. + request, _ := http.NewRequest(http.MethodPost, url, requestBody) + request.Header.Add("Content-Type", "application/json") + request.Header.Add("Accept", "application/json") + + resp, err := d.client.Do(request) + if err != nil { + return errors.Wrap(err, "doing snapshot field keys request") + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(resp.Body) + return errors.Errorf("status code: %d: %s", resp.StatusCode, b) + } + + return nil +} diff --git a/dax/mds/controller/naive/balancer.go b/dax/mds/controller/naive/balancer.go new file mode 100644 index 000000000..96644310b --- /dev/null +++ b/dax/mds/controller/naive/balancer.go @@ -0,0 +1,539 @@ +// Package naive contains a naive implementation of the Balancer interface. +package naive + +import ( + "context" + "fmt" + "math" + "sort" + "strings" + "sync" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/errors" + "github.com/molecula/featurebase/v3/logger" +) + +// Balancer is a naive implementation of the controller.Balancer interface. It +// helps manage the relationships between workers and jobs. The logic it uses to +// balance jobs across workers is very simple; it bases everything off the +// number of workers and number of jobs. It does not take anything else (such as +// job size, worker capabilities, etc) into consideration. +type Balancer struct { + mu sync.RWMutex + + // name is used in logging to help identify the balancer responsible for the + // log. + name string + + // current represents the current state of worker/job assigments. + current WorkerJobService + + // freeJobs is the set of jobs which have yet to be assigned to a worker. + // This could be because there are no available workers, or because a worker + // has been removed and the jobs for which it was responsible have yet to be + // reassigned. + freeJobs FreeJobService + + logger logger.Logger +} + +type WorkerJobService interface { + WorkersJobs(ctx context.Context, balancerName string) ([]dax.WorkerInfo, error) + + WorkerCount(ctx context.Context, balancerName string) (int, error) + ListWorkers(ctx context.Context, balancerName string) (dax.Workers, error) + WorkerExists(ctx context.Context, balancerName string, worker dax.Worker) (bool, error) + CreateWorker(ctx context.Context, balancerName string, worker dax.Worker) error + DeleteWorker(ctx context.Context, balancerName string, worker dax.Worker) error + + CreateJob(ctx context.Context, balancerName string, worker dax.Worker, job dax.Job) error + DeleteJob(ctx context.Context, balancerName string, worker dax.Worker, job dax.Job) error + JobCount(ctx context.Context, balancerName string, worker dax.Worker) (int, error) + ListJobs(ctx context.Context, balancerName string, worker dax.Worker) (dax.Jobs, error) +} + +type FreeJobService interface { + CreateFreeJob(ctx context.Context, balancerName string, job dax.Job) error + DeleteFreeJob(ctx context.Context, balancerName string, job dax.Job) error + ListFreeJobs(ctx context.Context, balancerName string) (dax.Jobs, error) + MergeFreeJobs(ctx context.Context, balancerName string, jobs dax.Jobs) error +} + +// New returns a new instance of Balancer. +func New(name string, fjs FreeJobService, wjs WorkerJobService, logger logger.Logger) *Balancer { + return &Balancer{ + name: name, + current: wjs, + freeJobs: fjs, + logger: logger, + } +} + +// AddWorker adds a worker to the Balancer's worker pool. This may cause the +// Balancer to assign existing jobs that are currently in the free list to the +// worker. Also, the worker will immediately be available for assignments of new +// jobs. +func (b *Balancer) AddWorker(ctx context.Context, worker fmt.Stringer) ([]dax.WorkerDiff, error) { + b.logger.Debugf("%s: AddWorker(%s)", b.name, worker.String()) + b.mu.Lock() + defer b.mu.Unlock() + + diff, err := b.addWorker(ctx, dax.Worker(worker.String())) + if err != nil { + return nil, errors.Wrap(err, "adding worker") + } + + return diff.output(), nil +} + +func (b *Balancer) addWorker(ctx context.Context, worker dax.Worker) (internalDiffs, error) { + // If this worker already exists, don't do anything. + if exists, err := b.current.WorkerExists(ctx, b.name, worker); err != nil { + return nil, errors.Wrap(err, "checking if worker exists") + } else if exists { + return internalDiffs{}, nil + } + + if err := b.current.CreateWorker(ctx, b.name, worker); err != nil { + return nil, errors.Wrap(err, "creating worker") + } + + // Process the freeJobs. + return b.processFreeJobs(ctx) +} + +// ReplaceWorker is meant to avoid the job re-assignment caused by performing a +// RemoveWorker followed by an AddWorker. In this case, it does both in one step +// so that it's more likely that the jobs will just get transferred directly +// over. NOT IMPLEMENTED YET. +// func (b *Balancer) ReplaceWorker(fromWorker string, toWorker string) []WorkerDiff { +// b.mu.Lock() +// defer b.mu.Unlock() + +// return []WorkerDiff{} +// } + +// RemoveWorker removes a worker from the worker pool and moves any of its +// currently assigned jobs to the free list. If the intention is to remove a +// worker and reassign its jobs to other workers, then RemoveWorker() should be +// followed by Balance(). +func (b *Balancer) RemoveWorker(ctx context.Context, worker fmt.Stringer) ([]dax.WorkerDiff, error) { + b.mu.Lock() + defer b.mu.Unlock() + + diff, err := b.removeWorker(ctx, dax.Worker(worker.String())) + if err != nil { + return nil, errors.Wrap(err, "removing worker") + } + + return diff.output(), nil +} + +func (b *Balancer) removeWorker(ctx context.Context, worker dax.Worker) (internalDiffs, error) { + // If this worker doesn't exist, don't do anything else. + if exists, err := b.current.WorkerExists(ctx, b.name, worker); err != nil { + return nil, errors.Wrap(err, "checking if worker exists") + } else if !exists { + return internalDiffs{}, nil + } + + jobs, err := b.current.ListJobs(ctx, b.name, worker) + if err != nil { + return nil, errors.Wrap(err, "listing jobs") + } + + // Before removing the worker, mark its jobs as free. + if err := b.freeJobs.MergeFreeJobs(ctx, b.name, jobs); err != nil { + return nil, errors.Wrap(err, "merging free jobs") + } + + // Remove the worker. + if err := b.current.DeleteWorker(ctx, b.name, worker); err != nil { + return nil, errors.Wrap(err, "deleting worker") + } + + // Even though this may not be useful to the caller (for example, in the + // case where the worker has died and no longer exists), return the diffs + // which represent the removal of jobs from the worker. + diff := newInternalDiffs() + for _, job := range jobs { + diff.removed(worker, job) + } + + return diff, nil +} + +// AddJob adds a job to an existing worker. If there are no existing workers, +// the job is placed into the free list and will be assigned to a worker once +// one becomes available. +func (b *Balancer) AddJob(ctx context.Context, job fmt.Stringer) ([]dax.WorkerDiff, error) { + b.logger.Debugf("%s: AddJob(%s)", b.name, job.String()) + b.mu.Lock() + defer b.mu.Unlock() + + diff, err := b.addJob(ctx, dax.Job(job.String())) + if err != nil { + return nil, errors.Wrap(err, "adding job") + } + + return diff.output(), nil +} + +func (b *Balancer) addJob(ctx context.Context, job dax.Job) (internalDiffs, error) { + if cnt, err := b.current.WorkerCount(ctx, b.name); err != nil { + return nil, errors.Wrap(err, "getting worker count") + } else if cnt == 0 { + if err := b.freeJobs.CreateFreeJob(ctx, b.name, job); err != nil { + return nil, errors.Wrap(err, "creating free job") + } + // TODO: we might want to inform the user that a job is in the free list + // because there are no workers. + return internalDiffs{}, nil + } + + // Make sure this job doesn't already exist. + if _, ok, err := b.workerForJob(ctx, job); err != nil { + return nil, errors.Wrapf(err, "getting worker for job: %s", job) + } else if ok { + // The job is already being tracked. + return internalDiffs{}, nil + } + + // Find the worker with the fewest number of jobs and assign it this job. + var lowCount int = math.MaxInt + var lowWorker dax.Worker + + workerIDs, err := b.current.ListWorkers(ctx, b.name) + if err != nil { + return nil, errors.Wrap(err, "listing workers") + } + + for _, workerID := range workerIDs { + if l, err := b.current.JobCount(ctx, b.name, workerID); err != nil { + return nil, errors.Wrapf(err, "getting job count for worker: %s", workerID) + } else if l < lowCount { + lowCount = l + lowWorker = workerID + } + } + + if err := b.current.CreateJob(ctx, b.name, lowWorker, job); err != nil { + return nil, errors.Wrap(err, "creating job") + } + + diffs := newInternalDiffs() + diffs.added(lowWorker, job) + + return diffs, nil +} + +// RemoveJob removes a job from the worker to which is was assigned. If the job +// is not currently assigned to a worker, but it is in the free list, then it +// will be removed from the free list. +func (b *Balancer) RemoveJob(ctx context.Context, job fmt.Stringer) ([]dax.WorkerDiff, error) { + b.mu.Lock() + defer b.mu.Unlock() + + diff, err := b.removeJob(ctx, dax.Job(job.String())) + if err != nil { + return nil, errors.Wrapf(err, "removing job: %s", job) + } + + return diff.output(), nil +} + +func (b *Balancer) removeJob(ctx context.Context, job dax.Job) (internalDiffs, error) { + if worker, ok, err := b.workerForJob(ctx, job); err != nil { + return nil, errors.Wrapf(err, "getting worker for job: %s", job) + } else if ok { + if err := b.current.DeleteJob(ctx, b.name, worker, job); err != nil { + return nil, errors.Wrapf(err, "deleting job: %s", job) + } + + diffs := newInternalDiffs() + diffs.removed(worker, job) + + return diffs, nil + } + + // Just in case the job is in the free list (and wasn't assigned to a + // worker), remove it; there's no need to provide a diff. There should never + // be a case where the same job is both in the free list and assigned to a + // worker. + if err := b.freeJobs.DeleteFreeJob(ctx, b.name, job); err != nil { + return nil, errors.Wrapf(err, "deleting free job: %s", job) + } + + return internalDiffs{}, nil +} + +// Balance ensures that all jobs are being handled by a worker by assigning jobs +// in the free list to workers, and by moving job assignments around in order to +// balance the load on workers. +func (b *Balancer) Balance(ctx context.Context) ([]dax.WorkerDiff, error) { + b.mu.Lock() + defer b.mu.Unlock() + + // If there are no workers, we can't properly balance. + if cnt, err := b.current.WorkerCount(ctx, b.name); err != nil { + return nil, errors.Wrapf(err, "getting worker count: %s", b.name) + } else if cnt == 0 { + return []dax.WorkerDiff{}, nil + } + + // Process the freeJobs. + diffs, err := b.processFreeJobs(ctx) + if err != nil { + return nil, errors.Wrapf(err, "processing free jobs: %s", b.name) + } + + // Balance the jobs among workers. + diff, err := b.balance(ctx, diffs) + if err != nil { + return nil, errors.Wrap(err, "balancing jobs") + } + + return diff.output(), nil +} + +// balance moves jobs among workers with the goal of having an equal number of +// jobs per worker. This method takes an `internalDiffs` as input for cases +// where some action has preceeded this call which also resulted in +// `internalDiffs`. Instead of having this method take a value, we could rely on +// the internalDiffs.merge() method, but we would need to modify that method to +// be smarter about the order in which it applies the add/remove operations. +// Until that's in place, we'll pass in a value here. +func (b *Balancer) balance(ctx context.Context, diffs internalDiffs) (internalDiffs, error) { + numWorkers, err := b.current.WorkerCount(ctx, b.name) + if err != nil { + return nil, errors.Wrapf(err, "getting worker count: %s", b.name) + } + numJobs := 0 + if workers, err := b.current.ListWorkers(ctx, b.name); err != nil { + return nil, errors.Wrapf(err, "listing workers: %s", b.name) + } else { + for _, worker := range workers { + cnt, err := b.current.JobCount(ctx, b.name, worker) + if err != nil { + return nil, errors.Wrapf(err, "getting job count: %s", worker) + } + numJobs += cnt + } + } + + minJobsPerWorker := numJobs / numWorkers + numWorkersAboveMin := numJobs % numWorkers + + // sortedWorkerInfos is used now in order to guarantee a sort order. + sortedWorkerInfos, err := b.currentState(ctx, true) + if err != nil { + return nil, errors.Wrapf(err, "getting current state: %s", b.name) + } + + // Loop through each worker, and if the number of jobs for the worker + // exceeds the target, then remove the job and add it back (which is + // effectively how we rebalance a job). + for i, workerInfo := range sortedWorkerInfos { + numTargetJobs := minJobsPerWorker + if i < numWorkersAboveMin { + numTargetJobs += 1 + } + + numCurrentJobs, err := b.current.JobCount(ctx, b.name, workerInfo.ID) + if err != nil { + return nil, errors.Wrapf(err, "getting job count: %s", workerInfo.ID) + } + + // If we don't need to remove jobs from this worker, then just continue + // on to the next worker. + if numCurrentJobs <= numTargetJobs { + continue + } + + sortedJobs, err := b.current.ListJobs(ctx, b.name, workerInfo.ID) + if err != nil { + return nil, errors.Wrapf(err, "listing jobs: %s", workerInfo.ID) + } + + // Remove the extra jobs from the end of the list, and add them back + // again (which should place them on a worker with fewer jobs). + for i := numCurrentJobs - 1; i >= numTargetJobs; i-- { + if rj, err := b.removeJob(ctx, sortedJobs[i]); err != nil { + return nil, errors.Wrapf(err, "removing job: %s", sortedJobs[i]) + } else { + diffs.merge(rj) + } + if aj, err := b.addJob(ctx, sortedJobs[i]); err != nil { + return nil, errors.Wrapf(err, "adding job: %s", sortedJobs[i]) + } else { + diffs.merge(aj) + } + } + } + + return diffs, nil +} + +// CurrentState returns the current state of worker and job assignments. Note +// that there could be unassigned jobs which are not captured in this output. +// Calling Balance() would force any unassigned jobs to be assigned (assuming +// there is at least one worker), and the output would then reflect that. +func (b *Balancer) CurrentState(ctx context.Context) ([]dax.WorkerInfo, error) { + b.mu.RLock() + defer b.mu.RUnlock() + + return b.currentState(ctx, true) +} + +func (b *Balancer) currentState(ctx context.Context, sorted bool) ([]dax.WorkerInfo, error) { + return b.current.WorkersJobs(ctx, b.name) +} + +// WorkerState returns the current state of job assignments for a given worker. +func (b *Balancer) WorkerState(ctx context.Context, worker dax.Worker) (dax.WorkerInfo, error) { + b.mu.RLock() + defer b.mu.RUnlock() + + return b.workerState(ctx, worker) +} + +func (b *Balancer) workerState(ctx context.Context, worker dax.Worker) (dax.WorkerInfo, error) { + if exists, err := b.current.WorkerExists(ctx, b.name, worker); err != nil { + return dax.WorkerInfo{}, errors.Wrapf(err, "checking worker exists: %s", worker) + } else if !exists { + return dax.WorkerInfo{ + ID: dax.Worker(worker), + }, nil + } + + jobs, err := b.current.ListJobs(ctx, b.name, worker) + if err != nil { + return dax.WorkerInfo{}, errors.Wrapf(err, "listing jobs: %s", worker) + } + + return dax.WorkerInfo{ + ID: dax.Worker(worker), + Jobs: jobs, + }, nil +} + +// WorkersForJobs returns the list of workers for the given jobs. If a given job +// is not currently assigned to a worker, it will be ignored. +func (b *Balancer) WorkersForJobs(ctx context.Context, jobs []dax.Job) ([]dax.WorkerInfo, error) { + b.mu.RLock() + defer b.mu.RUnlock() + + return b.workersForJobs(ctx, jobs) +} + +func (b *Balancer) workersForJobs(ctx context.Context, jobs []dax.Job) ([]dax.WorkerInfo, error) { + out := make(map[dax.Worker]dax.Set[dax.Job]) + + workerJobs, err := b.current.WorkersJobs(ctx, b.name) + if err != nil { + return nil, errors.Wrapf(err, "getting worker jobs: %s", b.name) + } + for _, workerInfo := range workerJobs { + jset := dax.NewSet(workerInfo.Jobs...) + + matches := dax.NewSet[dax.Job]() + for _, job := range jobs { + if jset.Contains(job) { + matches.Add(job) + } + } + + if len(matches) > 0 { + out[workerInfo.ID] = matches + } + } + + workers := make([]dax.WorkerInfo, len(out)) + + i := 0 + for w, jset := range out { + workers[i] = dax.WorkerInfo{ + ID: dax.Worker(w), + Jobs: jset.Sorted(), + } + i++ + } + + sort.Sort(dax.WorkerInfos(workers)) + + return workers, nil +} + +func (b *Balancer) WorkersForJobPrefix(ctx context.Context, prefix string) ([]dax.WorkerInfo, error) { + b.mu.RLock() + defer b.mu.RUnlock() + + jobs, err := b.freeJobs.ListFreeJobs(ctx, b.name) + if err != nil { + return nil, errors.Wrap(err, "listing free jobs") + } + for _, job := range jobs { + if strings.HasPrefix(string(job), prefix) { + return nil, errors.Errorf("found free job '%s' matching prefix '%s'", job, prefix) + } + } + + workerJobs, err := b.current.WorkersJobs(ctx, b.name) + if err != nil { + return nil, errors.Wrapf(err, "getting worker jobs: %s", b.name) + } + + result := make([]dax.WorkerInfo, 0) + for _, workerInfo := range workerJobs { + matchedJobs := make([]dax.Job, 0) + for _, job := range workerInfo.Jobs { + if strings.HasPrefix(string(job), prefix) { + matchedJobs = append(matchedJobs, job) + } + } + if len(matchedJobs) > 0 { + result = append(result, dax.WorkerInfo{ + ID: workerInfo.ID, + Jobs: matchedJobs, + }) + } + } + return result, nil + +} + +// processFreeJobs assigns all jobs in the free list to a worker. +func (b *Balancer) processFreeJobs(ctx context.Context) (internalDiffs, error) { + diffs := newInternalDiffs() + jobs, err := b.freeJobs.ListFreeJobs(ctx, b.name) + if err != nil { + return nil, errors.Wrapf(err, "listing free jobs: %s", b.name) + } + for _, job := range jobs { + if aj, err := b.addJob(ctx, job); err != nil { + return nil, errors.Wrapf(err, "adding job: %s", job) + } else { + diffs.merge(aj) + } + if err := b.freeJobs.DeleteFreeJob(ctx, b.name, job); err != nil { + return nil, errors.Wrapf(err, "deleting free job: %s", job) + } + } + return diffs, nil +} + +// workerForJob returns the worker currently assigned to the given job. +func (b *Balancer) workerForJob(ctx context.Context, job dax.Job) (dax.Worker, bool, error) { + workerJobs, err := b.current.WorkersJobs(ctx, b.name) + if err != nil { + return "", false, errors.Wrapf(err, "getting workers jobs: %s", b.name) + } + for _, workerInfo := range workerJobs { + jset := dax.NewSet(workerInfo.Jobs...) + if jset.Contains(job) { + return workerInfo.ID, true, nil + } + } + return "", false, nil +} diff --git a/dax/mds/controller/naive/balancer_test.go b/dax/mds/controller/naive/balancer_test.go new file mode 100644 index 000000000..f4f9a8095 --- /dev/null +++ b/dax/mds/controller/naive/balancer_test.go @@ -0,0 +1,739 @@ +package naive_test + +import ( + "context" + "fmt" + "os" + "testing" + + "github.com/molecula/featurebase/v3/dax" + daxbolt "github.com/molecula/featurebase/v3/dax/boltdb" + "github.com/molecula/featurebase/v3/dax/mds/controller/naive/boltdb" + testbolt "github.com/molecula/featurebase/v3/dax/test/boltdb" + "github.com/molecula/featurebase/v3/logger" + "github.com/stretchr/testify/assert" +) + +func newBoltBalancer(t *testing.T) (*daxbolt.DB, func()) { + db := testbolt.MustOpenDB(t) + assert.NoError(t, db.InitializeBuckets(boltdb.NaiveBalancerBuckets...)) + + return db, func() { + testbolt.MustCloseDB(t, db) + testbolt.CleanupDB(t, db.Path()) + } +} + +func TestBalancer(t *testing.T) { + ctx := context.Background() + t.Run("SingleWorker", func(t *testing.T) { + db, cleanup := newBoltBalancer(t) + defer cleanup() + bal := boltdb.NewBalancer("test", db, logger.NewStandardLogger(os.Stderr)) + tests := []struct { + fn func(context.Context, fmt.Stringer) ([]dax.WorkerDiff, error) + input string + expDiff []dax.WorkerDiff + expState []dax.WorkerInfo + }{ + { + // Add job. + fn: bal.AddJob, + input: "p2", + expDiff: []dax.WorkerDiff{}, + expState: []dax.WorkerInfo{}, + }, + { + // Add worker. + fn: bal.AddWorker, + input: "n1", + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n1", + AddedJobs: []dax.Job{"p2"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n1", + Jobs: []dax.Job{"p2"}, + }, + }, + }, + { + // Add another job out of order. + fn: bal.AddJob, + input: "p1", + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n1", + AddedJobs: []dax.Job{"p1"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n1", + Jobs: []dax.Job{"p1", "p2"}, + }, + }, + }, + { + // Add another job. + fn: bal.AddJob, + input: "p3", + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n1", + AddedJobs: []dax.Job{"p3"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n1", + Jobs: []dax.Job{"p1", "p2", "p3"}, + }, + }, + }, + { + // Add a duplicate job. + fn: bal.AddJob, + input: "p2", + expDiff: []dax.WorkerDiff{}, + expState: []dax.WorkerInfo{ + { + ID: "n1", + Jobs: []dax.Job{"p1", "p2", "p3"}, + }, + }, + }, + } + for i, test := range tests { + t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) { + diff, err := test.fn(ctx, newStringWrapper(test.input)) + assert.NoError(t, err) + assert.Equal(t, test.expDiff, diff) + + cs, err := bal.CurrentState(ctx) + assert.NoError(t, err) + assert.Equal(t, test.expState, cs) + }) + } + }) + + t.Run("MultipleWorkers", func(t *testing.T) { + db, cleanup := newBoltBalancer(t) + defer cleanup() + bal := boltdb.NewBalancer("test", db, logger.NewStandardLogger(os.Stderr)) + tests := []struct { + fn func(context.Context, fmt.Stringer) ([]dax.WorkerDiff, error) + input string + balance bool + expDiff []dax.WorkerDiff + expState []dax.WorkerInfo + }{ + { + // Balance when empty. + balance: true, + expDiff: []dax.WorkerDiff{}, + expState: []dax.WorkerInfo{}, + }, + { + // Add worker. + fn: bal.AddWorker, + input: "n2", + expDiff: []dax.WorkerDiff{}, + expState: []dax.WorkerInfo{ + { + ID: "n2", + Jobs: []dax.Job{}, + }, + }, + }, + { + // Add worker again. + fn: bal.AddWorker, + input: "n2", + expDiff: []dax.WorkerDiff{}, + expState: []dax.WorkerInfo{ + { + ID: "n2", + Jobs: []dax.Job{}, + }, + }, + }, + { + // Add a second worker. + fn: bal.AddWorker, + input: "n1", + expDiff: []dax.WorkerDiff{}, + expState: []dax.WorkerInfo{ + { + ID: "n1", + Jobs: []dax.Job{}, + }, + { + ID: "n2", + Jobs: []dax.Job{}, + }, + }, + }, + { + // Add job. + fn: bal.AddJob, + input: "p2", + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n1", + AddedJobs: []dax.Job{"p2"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n1", + Jobs: []dax.Job{"p2"}, + }, + { + ID: "n2", + Jobs: []dax.Job{}, + }, + }, + }, + { + // Add job. + fn: bal.AddJob, + input: "p3", + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n2", + AddedJobs: []dax.Job{"p3"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n1", + Jobs: []dax.Job{"p2"}, + }, + { + ID: "n2", + Jobs: []dax.Job{"p3"}, + }, + }, + }, + { + // Add job. + fn: bal.AddJob, + input: "p1", + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n1", + AddedJobs: []dax.Job{"p1"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n1", + Jobs: []dax.Job{"p1", "p2"}, + }, + { + ID: "n2", + Jobs: []dax.Job{"p3"}, + }, + }, + }, + { + // Add a third worker. + fn: bal.AddWorker, + input: "n0", + expDiff: []dax.WorkerDiff{}, + expState: []dax.WorkerInfo{ + { + ID: "n0", + Jobs: []dax.Job{}, + }, + { + ID: "n1", + Jobs: []dax.Job{"p1", "p2"}, + }, + { + ID: "n2", + Jobs: []dax.Job{"p3"}, + }, + }, + }, + { + // Add job. + fn: bal.AddJob, + input: "p4", + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n0", + AddedJobs: []dax.Job{"p4"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n0", + Jobs: []dax.Job{"p4"}, + }, + { + ID: "n1", + Jobs: []dax.Job{"p1", "p2"}, + }, + { + ID: "n2", + Jobs: []dax.Job{"p3"}, + }, + }, + }, + { + // Add job. + fn: bal.AddJob, + input: "p5", + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n0", + AddedJobs: []dax.Job{"p5"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n0", + Jobs: []dax.Job{"p4", "p5"}, + }, + { + ID: "n1", + Jobs: []dax.Job{"p1", "p2"}, + }, + { + ID: "n2", + Jobs: []dax.Job{"p3"}, + }, + }, + }, + { + // Add job. + fn: bal.AddJob, + input: "p0", + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n2", + AddedJobs: []dax.Job{"p0"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n0", + Jobs: []dax.Job{"p4", "p5"}, + }, + { + ID: "n1", + Jobs: []dax.Job{"p1", "p2"}, + }, + { + ID: "n2", + Jobs: []dax.Job{"p0", "p3"}, + }, + }, + }, + { + // Add job. + fn: bal.AddJob, + input: "p6", + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n0", + AddedJobs: []dax.Job{"p6"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n0", + Jobs: []dax.Job{"p4", "p5", "p6"}, + }, + { + ID: "n1", + Jobs: []dax.Job{"p1", "p2"}, + }, + { + ID: "n2", + Jobs: []dax.Job{"p0", "p3"}, + }, + }, + }, + { + // Add job. + fn: bal.AddJob, + input: "p7", + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n1", + AddedJobs: []dax.Job{"p7"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n0", + Jobs: []dax.Job{"p4", "p5", "p6"}, + }, + { + ID: "n1", + Jobs: []dax.Job{"p1", "p2", "p7"}, + }, + { + ID: "n2", + Jobs: []dax.Job{"p0", "p3"}, + }, + }, + }, + + //////////////////// Remove ///////////////////////// + + { + // Remove nonexistent worker. + fn: bal.RemoveWorker, + input: "nonexistent", + expDiff: []dax.WorkerDiff{}, + expState: []dax.WorkerInfo{ + { + ID: "n0", + Jobs: []dax.Job{"p4", "p5", "p6"}, + }, + { + ID: "n1", + Jobs: []dax.Job{"p1", "p2", "p7"}, + }, + { + ID: "n2", + Jobs: []dax.Job{"p0", "p3"}, + }, + }, + }, + { + // Remove worker. + fn: bal.RemoveWorker, + input: "n1", + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n1", + AddedJobs: []dax.Job{}, + RemovedJobs: []dax.Job{"p1", "p2", "p7"}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n0", + Jobs: []dax.Job{"p4", "p5", "p6"}, + }, + { + ID: "n2", + Jobs: []dax.Job{"p0", "p3"}, + }, + }, + }, + + { + // Remove job (from free list). + fn: bal.RemoveJob, + input: "p2", + expDiff: []dax.WorkerDiff{}, + expState: []dax.WorkerInfo{ + { + ID: "n0", + Jobs: []dax.Job{"p4", "p5", "p6"}, + }, + { + ID: "n2", + Jobs: []dax.Job{"p0", "p3"}, + }, + }, + }, + + { + // Balance after remove. + balance: true, + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n0", + AddedJobs: []dax.Job{"p7"}, + RemovedJobs: []dax.Job{}, + }, + { + WorkerID: "n2", + AddedJobs: []dax.Job{"p1"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n0", + Jobs: []dax.Job{"p4", "p5", "p6", "p7"}, + }, + { + ID: "n2", + Jobs: []dax.Job{"p0", "p1", "p3"}, + }, + }, + }, + + { + // Remove job. + fn: bal.RemoveJob, + input: "p1", + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n2", + AddedJobs: []dax.Job{}, + RemovedJobs: []dax.Job{"p1"}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n0", + Jobs: []dax.Job{"p4", "p5", "p6", "p7"}, + }, + { + ID: "n2", + Jobs: []dax.Job{"p0", "p3"}, + }, + }, + }, + } + for i, test := range tests { + t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) { + var diff []dax.WorkerDiff + var err error + if test.balance { + diff, err = bal.Balance(ctx) + } else { + diff, err = test.fn(ctx, newStringWrapper(test.input)) + } + assert.NoError(t, err) + assert.Equal(t, test.expDiff, diff) + + cs, err := bal.CurrentState(ctx) + assert.NoError(t, err) + assert.Equal(t, test.expState, cs) + }) + } + }) + + t.Run("WorkerState", func(t *testing.T) { + db, cleanup := newBoltBalancer(t) + defer cleanup() + bal := boltdb.NewBalancer("test", db, logger.NewStandardLogger(os.Stderr)) + + _, err := bal.AddWorker(ctx, newStringWrapper("n1")) + assert.NoError(t, err) + _, err = bal.AddJob(ctx, newStringWrapper("p1")) + assert.NoError(t, err) + + exp := dax.WorkerInfo{ + ID: "n1", + Jobs: []dax.Job{"p1"}, + } + ws, err := bal.WorkerState(ctx, "n1") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + + // Worker doesn't exist. + exp = dax.WorkerInfo{ + ID: "x1", + } + ws, err = bal.WorkerState(ctx, "x1") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + }) + + t.Run("WorkersForJobs", func(t *testing.T) { + db, cleanup := newBoltBalancer(t) + defer cleanup() + bal := boltdb.NewBalancer("test", db, logger.NewStandardLogger(os.Stderr)) + + _, err := bal.AddWorker(ctx, newStringWrapper("n1")) + assert.NoError(t, err) + _, err = bal.AddWorker(ctx, newStringWrapper("n2")) + assert.NoError(t, err) + for i := 0; i < 12; i++ { + _, err = bal.AddJob(ctx, newStringWrapper(fmt.Sprintf("p%d", i))) + assert.NoError(t, err) + } + + exp := dax.WorkerInfo{ + ID: "n1", + Jobs: []dax.Job{"p0", "p10", "p2", "p4", "p6", "p8"}, + } + ws, err := bal.WorkerState(ctx, "n1") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + + exp = dax.WorkerInfo{ + ID: "n2", + Jobs: []dax.Job{"p1", "p11", "p3", "p5", "p7", "p9"}, + } + ws, err = bal.WorkerState(ctx, "n2") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + + tests := []struct { + jobs []dax.Job + exp []dax.WorkerInfo + }{ + { + jobs: []dax.Job{"p0"}, + exp: []dax.WorkerInfo{ + {ID: "n1", Jobs: []dax.Job{"p0"}}, + }, + }, + { + jobs: []dax.Job{"p0", "p4"}, + exp: []dax.WorkerInfo{ + {ID: "n1", Jobs: []dax.Job{"p0", "p4"}}, + }, + }, + { + jobs: []dax.Job{"p0", "p4", "p999"}, + exp: []dax.WorkerInfo{ + {ID: "n1", Jobs: []dax.Job{"p0", "p4"}}, + }, + }, + { + jobs: []dax.Job{"p0", "p1"}, + exp: []dax.WorkerInfo{ + {ID: "n1", Jobs: []dax.Job{"p0"}}, + {ID: "n2", Jobs: []dax.Job{"p1"}}, + }, + }, + { + jobs: []dax.Job{"p5", "p0", "p1", "p8"}, + exp: []dax.WorkerInfo{ + {ID: "n1", Jobs: []dax.Job{"p0", "p8"}}, + {ID: "n2", Jobs: []dax.Job{"p1", "p5"}}, + }, + }, + } + for i, test := range tests { + t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) { + workers, err := bal.WorkersForJobs(ctx, test.jobs) + assert.NoError(t, err) + assert.Equal(t, test.exp, workers) + }) + } + + // Some tests for WorkersForJobPrefix + workers, err := bal.WorkersForJobPrefix(ctx, "p1") + assert.NoError(t, err) + assert.ElementsMatch(t, []dax.WorkerInfo{ + {ID: "n1", Jobs: []dax.Job{"p10"}}, + {ID: "n2", Jobs: []dax.Job{"p1", "p11"}}, + }, workers) + + workers, err = bal.WorkersForJobPrefix(ctx, "p2") + assert.NoError(t, err) + assert.ElementsMatch(t, []dax.WorkerInfo{ + {ID: "n1", Jobs: []dax.Job{"p2"}}, + }, workers) + + workers, err = bal.WorkersForJobPrefix(ctx, "pp") + assert.NoError(t, err) + assert.ElementsMatch(t, []dax.WorkerInfo{}, workers) + + }) + + t.Run("Balance", func(t *testing.T) { + db, cleanup := newBoltBalancer(t) + defer cleanup() + bal := boltdb.NewBalancer("test", db, logger.NewStandardLogger(os.Stderr)) + + // Add two workers with some jobs evenly spread across them. + _, err := bal.AddWorker(ctx, newStringWrapper("n1")) + assert.NoError(t, err) + _, err = bal.AddWorker(ctx, newStringWrapper("n2")) + assert.NoError(t, err) + for i := 0; i < 13; i++ { + _, err = bal.AddJob(ctx, newStringWrapper(fmt.Sprintf("p%d", i))) + assert.NoError(t, err) + } + + exp := dax.WorkerInfo{ + ID: "n1", + Jobs: []dax.Job{"p0", "p10", "p12", "p2", "p4", "p6", "p8"}, + } + ws, err := bal.WorkerState(ctx, "n1") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + + exp = dax.WorkerInfo{ + ID: "n2", + Jobs: []dax.Job{"p1", "p11", "p3", "p5", "p7", "p9"}, + } + ws, err = bal.WorkerState(ctx, "n2") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + + // Now, add a worker and confirm that it currently has no jobs assigned + // to it. + _, err = bal.AddWorker(ctx, newStringWrapper("n3")) + assert.NoError(t, err) + exp = dax.WorkerInfo{ + ID: "n3", + Jobs: []dax.Job{}, + } + ws, err = bal.WorkerState(ctx, "n3") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + + // Finally, call Balance() and confirm that the appropriate jobs got + // reassigned. + _, err = bal.Balance(ctx) + assert.NoError(t, err) + + exp = dax.WorkerInfo{ + ID: "n1", + Jobs: []dax.Job{"p0", "p10", "p12", "p2", "p4"}, + } + ws, err = bal.WorkerState(ctx, "n1") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + + exp = dax.WorkerInfo{ + ID: "n2", + Jobs: []dax.Job{"p1", "p11", "p3", "p5"}, + } + ws, err = bal.WorkerState(ctx, "n2") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + + exp = dax.WorkerInfo{ + ID: "n3", + Jobs: []dax.Job{"p6", "p7", "p8", "p9"}, + } + ws, err = bal.WorkerState(ctx, "n3") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + }) +} + +type stringWrapper struct { + s string +} + +func newStringWrapper(s string) *stringWrapper { + return &stringWrapper{ + s: s, + } +} + +func (s *stringWrapper) String() string { + return s.s +} diff --git a/dax/mds/controller/naive/boltdb/balancer.go b/dax/mds/controller/naive/boltdb/balancer.go new file mode 100644 index 000000000..d943eba6d --- /dev/null +++ b/dax/mds/controller/naive/boltdb/balancer.go @@ -0,0 +1,515 @@ +// Package boltdb contains the boltdb implementation of the Balancer interface. +package boltdb + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "strings" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/boltdb" + "github.com/molecula/featurebase/v3/dax/mds/controller" + "github.com/molecula/featurebase/v3/dax/mds/controller/naive" + "github.com/molecula/featurebase/v3/errors" + "github.com/molecula/featurebase/v3/logger" +) + +var ( + bucketNaiveBalancer = boltdb.Bucket("naiveBalancer") +) + +// NaiveBalancerBuckets defines the buckets used by this package. It can be +// called during setup to create the buckets ahead of time. +var NaiveBalancerBuckets []boltdb.Bucket = []boltdb.Bucket{ + bucketNaiveBalancer, +} + +// NewBalancer returns a new instance of controller.Balancer. +func NewBalancer(name string, db *boltdb.DB, logger logger.Logger) controller.Balancer { + fjs := newFreeJobService(db) + wjs := newWorkerJobService(db, logger) + + return naive.New(name, fjs, wjs, logger) +} + +// Ensure type implements interface. +var _ naive.WorkerJobService = (*workerJobService)(nil) + +type workerJobService struct { + db *boltdb.DB + logger logger.Logger +} + +func newWorkerJobService(db *boltdb.DB, logger logger.Logger) *workerJobService { + return &workerJobService{ + db: db, + logger: logger, + } +} + +func (w *workerJobService) WorkersJobs(ctx context.Context, balancerName string) ([]dax.WorkerInfo, error) { + tx, err := w.db.BeginTx(ctx, false) + if err != nil { + return nil, errors.Wrap(err, "getting tx") + } + defer tx.Rollback() + + workerInfos, err := getWorkerInfos(ctx, tx, balancerName) + if err != nil { + return nil, errors.Wrapf(err, "getting worker infos: %s", balancerName) + } + + return workerInfos, nil +} + +func (w *workerJobService) WorkerCount(ctx context.Context, balancerName string) (int, error) { + tx, err := w.db.BeginTx(ctx, false) + if err != nil { + return 0, errors.Wrap(err, "getting tx") + } + defer tx.Rollback() + + workers, err := w.getWorkers(ctx, tx, balancerName) + if err != nil { + return 0, errors.Wrapf(err, "getting workers: %s", balancerName) + } + + return len(workers), nil +} + +func (w *workerJobService) ListWorkers(ctx context.Context, balancerName string) (dax.Workers, error) { + tx, err := w.db.BeginTx(ctx, false) + if err != nil { + return nil, errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + workers, err := w.getWorkers(ctx, tx, balancerName) + if err != nil { + return nil, errors.Wrapf(err, "getting workers: %s", balancerName) + } + + return workers, nil +} + +func (w *workerJobService) getWorkers(ctx context.Context, tx *boltdb.Tx, balancerName string) (dax.Workers, error) { + c := tx.Bucket(bucketNaiveBalancer).Cursor() + + // Deserialize rows into Worker objects. + workers := make(dax.Workers, 0) + + prefix := []byte(fmt.Sprintf(prefixFmtWorkers, balancerName)) + for k, v := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, v = c.Next() { + if v == nil { + w.logger.Printf("nil value for key: %s", k) + continue + } + + worker, err := keyWorker(k) + if err != nil { + return nil, errors.Wrapf(err, "getting worker from key: %v", k) + } + + workers = append(workers, worker) + } + + return workers, nil +} + +func getWorkerInfos(ctx context.Context, tx *boltdb.Tx, balancerName string) (dax.WorkerInfos, error) { + c := tx.Bucket(bucketNaiveBalancer).Cursor() + + // Deserialize rows into WorkerInfo objects. + workerInfos := make(dax.WorkerInfos, 0) + + prefix := []byte(fmt.Sprintf(prefixFmtWorkers, balancerName)) + for k, v := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, v = c.Next() { + worker, err := keyWorker(k) + if err != nil { + return nil, errors.Wrapf(err, "getting worker from key: %v", k) + } + + jobs := dax.NewSet[dax.Job]() + if v != nil { + jobs, err = decodeJobSet(v) + if err != nil { + return nil, errors.Wrap(err, "decoding job set") + } + } + + workerInfo := dax.WorkerInfo{ + ID: worker, + Jobs: jobs.Sorted(), + } + + workerInfos = append(workerInfos, workerInfo) + } + + return workerInfos, nil +} + +func (w *workerJobService) WorkerExists(ctx context.Context, balancerName string, worker dax.Worker) (bool, error) { + tx, err := w.db.BeginTx(ctx, false) + if err != nil { + return false, errors.Wrapf(err, "getting tx: %s", balancerName) + } + defer tx.Rollback() + + bkt := tx.Bucket(bucketNaiveBalancer) + if bkt == nil { + return false, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer) + } + + wrkr := bkt.Get(workerKey(balancerName, worker)) + + return wrkr != nil, nil +} + +func (w *workerJobService) CreateWorker(ctx context.Context, balancerName string, worker dax.Worker) error { + tx, err := w.db.BeginTx(ctx, true) + if err != nil { + return errors.Wrap(err, "getting transaction") + } + defer tx.Rollback() + + bkt := tx.Bucket(bucketNaiveBalancer) + if bkt == nil { + return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer) + } + + // If this worker already exists, don't do anything. + wrkr := bkt.Get(workerKey(balancerName, worker)) + if wrkr != nil { + return nil + } + + val := []byte("[]") + if err := bkt.Put(workerKey(balancerName, worker), val); err != nil { + return errors.Wrap(err, "putting worker") + } + + return tx.Commit() +} + +func (w *workerJobService) DeleteWorker(ctx context.Context, balancerName string, worker dax.Worker) error { + tx, err := w.db.BeginTx(ctx, true) + if err != nil { + return errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + bkt := tx.Bucket(bucketNaiveBalancer) + if bkt == nil { + return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer) + } + + if err := bkt.Delete(workerKey(balancerName, worker)); err != nil { + return errors.Wrapf(err, "deleting node key: %s", workerKey(balancerName, worker)) + } + + return tx.Commit() +} + +func (w *workerJobService) CreateJob(ctx context.Context, balancerName string, worker dax.Worker, job dax.Job) error { + tx, err := w.db.BeginTx(ctx, true) + if err != nil { + return errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + bkt := tx.Bucket(bucketNaiveBalancer) + if bkt == nil { + return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer) + } + + jobset := dax.NewSet[dax.Job]() + + // get worker + wrkr := bkt.Get(workerKey(balancerName, worker)) + if wrkr != nil { + jobset, err = decodeJobSet(wrkr) + if err != nil { + return errors.Wrap(err, "decoding job set") + } + } + + jobset.Add(job) + val, err := encodeJobSet(jobset) + if err != nil { + return errors.Wrap(err, "encoding job set") + } + + if err := bkt.Put(workerKey(balancerName, worker), val); err != nil { + return errors.Wrap(err, "putting worker") + } + + return tx.Commit() +} + +func (w *workerJobService) DeleteJob(ctx context.Context, balancerName string, worker dax.Worker, job dax.Job) error { + tx, err := w.db.BeginTx(ctx, true) + if err != nil { + return errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + bkt := tx.Bucket(bucketNaiveBalancer) + if bkt == nil { + return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer) + } + + // get worker + wrkr := bkt.Get(workerKey(balancerName, worker)) + if wrkr == nil { + return nil + } + + jobset, err := decodeJobSet(wrkr) + if err != nil { + return errors.Wrap(err, "decoding job set") + } + if !jobset.Contains(job) { + return nil + } + + jobset.Remove(job) + val, err := encodeJobSet(jobset) + if err != nil { + return errors.Wrap(err, "encoding job set") + } + + if err := bkt.Put(workerKey(balancerName, worker), val); err != nil { + return errors.Wrap(err, "putting worker") + } + + return tx.Commit() +} + +func (w *workerJobService) ListJobs(ctx context.Context, balancerName string, worker dax.Worker) (dax.Jobs, error) { + tx, err := w.db.BeginTx(ctx, false) + if err != nil { + return nil, errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + bkt := tx.Bucket(bucketNaiveBalancer) + if bkt == nil { + return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer) + } + + jobset := dax.NewSet[dax.Job]() + + // get worker + wrkr := bkt.Get(workerKey(balancerName, worker)) + if wrkr != nil { + jobset, err = decodeJobSet(wrkr) + if err != nil { + return nil, errors.Wrap(err, "decoding job set") + } + } + + return jobset.Sorted(), nil +} + +func (w *workerJobService) JobCount(ctx context.Context, balancerName string, worker dax.Worker) (int, error) { + tx, err := w.db.BeginTx(ctx, false) + if err != nil { + return 0, errors.Wrapf(err, "getting tx: %s", balancerName) + } + defer tx.Rollback() + + bkt := tx.Bucket(bucketNaiveBalancer) + if bkt == nil { + return 0, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer) + } + + jobset := dax.NewSet[dax.Job]() + + // get worker + wrkr := bkt.Get(workerKey(balancerName, worker)) + if wrkr != nil { + jobset, err = decodeJobSet(wrkr) + if err != nil { + return 0, errors.Wrap(err, "decoding job set") + } + } + + return len(jobset), nil +} + +// encodeJobSet encode the jobSet into a JSON array of strings. +func encodeJobSet(jobSet dax.Set[dax.Job]) ([]byte, error) { + arr := jobSet.Sorted() + b, err := json.Marshal(arr) + if err != nil { + return nil, errors.Wrap(err, "marshalling json") + } + return b, nil +} + +// decodeJobSet decode the string (a JSON array of strings) into jobSet. +func decodeJobSet(v []byte) (dax.Set[dax.Job], error) { + var arr []string + err := json.Unmarshal(v, &arr) + if err != nil { + return nil, errors.Wrap(err, "unmarshalling json") + } + + js := dax.NewSet[dax.Job]() + for _, s := range arr { + js.Add(dax.Job(s)) + } + + return js, nil +} + +// Ensure type implements interface. +var _ naive.FreeJobService = (*freeJobService)(nil) + +type freeJobService struct { + db *boltdb.DB +} + +func newFreeJobService(db *boltdb.DB) *freeJobService { + return &freeJobService{ + db: db, + } +} + +func (f *freeJobService) CreateFreeJob(ctx context.Context, balancerName string, job dax.Job) error { + return f.MergeFreeJobs(ctx, balancerName, dax.Jobs{job}) +} + +func (f *freeJobService) DeleteFreeJob(ctx context.Context, balancerName string, job dax.Job) error { + tx, err := f.db.BeginTx(ctx, true) + if err != nil { + return errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + bkt := tx.Bucket(bucketNaiveBalancer) + if bkt == nil { + return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer) + } + + // get free jobs + fjs := bkt.Get(freeJobKey(balancerName)) + if fjs == nil { + return nil + } + + jobset, err := decodeJobSet(fjs) + if err != nil { + return errors.Wrap(err, "decoding job set") + } + if !jobset.Contains(job) { + return nil + } + + jobset.Remove(job) + val, err := encodeJobSet(jobset) + if err != nil { + return errors.Wrap(err, "encoding job set") + } + + if err := bkt.Put(freeJobKey(balancerName), val); err != nil { + return errors.Wrap(err, "putting free job") + } + + return tx.Commit() +} + +func (f *freeJobService) ListFreeJobs(ctx context.Context, balancerName string) (dax.Jobs, error) { + tx, err := f.db.BeginTx(ctx, false) + if err != nil { + return nil, errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + bkt := tx.Bucket(bucketNaiveBalancer) + if bkt == nil { + return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer) + } + + jobset := dax.NewSet[dax.Job]() + + // get free jobs + fjs := bkt.Get(freeJobKey(balancerName)) + if fjs != nil { + jobset, err = decodeJobSet(fjs) + if err != nil { + return nil, errors.Wrap(err, "decoding job set") + } + } + + return jobset.Sorted(), nil +} + +func (f *freeJobService) MergeFreeJobs(ctx context.Context, balancerName string, jobs dax.Jobs) error { + tx, err := f.db.BeginTx(ctx, true) + if err != nil { + return err + } + defer tx.Rollback() + + bkt := tx.Bucket(bucketNaiveBalancer) + if bkt == nil { + return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketNaiveBalancer) + } + + jobset := dax.NewSet[dax.Job]() + + // get free jobs + fjs := bkt.Get(freeJobKey(balancerName)) + if fjs != nil { + jobset, err = decodeJobSet(fjs) + if err != nil { + return errors.Wrap(err, "decoding job set") + } + } + + for _, j := range jobs { + jobset.Add(j) + } + val, err := encodeJobSet(jobset) + if err != nil { + return errors.Wrap(err, "encoding job set") + } + + if err := bkt.Put(freeJobKey(balancerName), val); err != nil { + return errors.Wrap(err, "putting free job") + } + + return tx.Commit() +} + +////////////////////////////////////////////////////// + +const ( + prefixFmtWorkers = "workers/%s/" // %s - balancerName + prefixFmtFreeJobs = "freejobs/%s" // %s - balancerName +) + +// workerKey returns a key based on worker. +func workerKey(bal string, worker dax.Worker) []byte { + key := fmt.Sprintf(prefixFmtWorkers+"%s", bal, worker) + return []byte(key) +} + +// keyWorker gets the worker out of the key. +func keyWorker(key []byte) (dax.Worker, error) { + parts := strings.Split(string(key), "/") + if len(parts) != 3 { + return "", errors.New(errors.ErrUncoded, "worker key format expected: `workers/balancer/worker`") + } + + return dax.Worker(parts[2]), nil +} + +// freeJobKey returns a key for all freeJobs. +func freeJobKey(bal string) []byte { + key := fmt.Sprintf(prefixFmtFreeJobs, bal) + return []byte(key) +} diff --git a/dax/mds/controller/naive/boltdb/balancer_test.go b/dax/mds/controller/naive/boltdb/balancer_test.go new file mode 100644 index 000000000..ebc28856d --- /dev/null +++ b/dax/mds/controller/naive/boltdb/balancer_test.go @@ -0,0 +1,709 @@ +package boltdb_test + +import ( + "context" + "fmt" + "testing" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/mds/controller/naive/boltdb" + testbolt "github.com/molecula/featurebase/v3/dax/test/boltdb" + "github.com/molecula/featurebase/v3/logger" + "github.com/stretchr/testify/assert" +) + +func TestBalancer(t *testing.T) { + db := testbolt.MustOpenDB(t) + defer testbolt.MustCloseDB(t, db) + + t.Cleanup(func() { + testbolt.CleanupDB(t, db.Path()) + }) + + ctx := context.Background() + + // Initialize the buckets. + assert.NoError(t, db.InitializeBuckets(boltdb.NaiveBalancerBuckets...)) + + t.Run("SingleWorker", func(t *testing.T) { + bal := boltdb.NewBalancer("test-single-worker", db, logger.NopLogger) + tests := []struct { + fn func(context.Context, fmt.Stringer) ([]dax.WorkerDiff, error) + input string + expDiff []dax.WorkerDiff + expState []dax.WorkerInfo + }{ + { + // Add job. + fn: bal.AddJob, + input: "p2", + expDiff: []dax.WorkerDiff{}, + expState: []dax.WorkerInfo{}, + }, + { + // Add worker. + fn: bal.AddWorker, + input: "n1", + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n1", + AddedJobs: []dax.Job{"p2"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n1", + Jobs: []dax.Job{"p2"}, + }, + }, + }, + { + // Add another job out of order. + fn: bal.AddJob, + input: "p1", + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n1", + AddedJobs: []dax.Job{"p1"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n1", + Jobs: []dax.Job{"p1", "p2"}, + }, + }, + }, + { + // Add another job. + fn: bal.AddJob, + input: "p3", + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n1", + AddedJobs: []dax.Job{"p3"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n1", + Jobs: []dax.Job{"p1", "p2", "p3"}, + }, + }, + }, + { + // Add a duplicate job. + fn: bal.AddJob, + input: "p2", + expDiff: []dax.WorkerDiff{}, + expState: []dax.WorkerInfo{ + { + ID: "n1", + Jobs: []dax.Job{"p1", "p2", "p3"}, + }, + }, + }, + } + for i, test := range tests { + t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) { + diff, err := test.fn(ctx, newStringWrapper(test.input)) + assert.NoError(t, err) + assert.Equal(t, test.expDiff, diff) + + cs, err := bal.CurrentState(ctx) + assert.NoError(t, err) + assert.Equal(t, test.expState, cs) + }) + } + }) + + t.Run("MultipleWorkers", func(t *testing.T) { + bal := boltdb.NewBalancer("test-multiple-workers", db, logger.NopLogger) + tests := []struct { + fn func(context.Context, fmt.Stringer) ([]dax.WorkerDiff, error) + input string + balance bool + expDiff []dax.WorkerDiff + expState []dax.WorkerInfo + }{ + { + // Balance when empty. + balance: true, + expDiff: []dax.WorkerDiff{}, + expState: []dax.WorkerInfo{}, + }, + { + // Add worker. + fn: bal.AddWorker, + input: "n2", + expDiff: []dax.WorkerDiff{}, + expState: []dax.WorkerInfo{ + { + ID: "n2", + Jobs: []dax.Job{}, + }, + }, + }, + { + // Add worker again. + fn: bal.AddWorker, + input: "n2", + expDiff: []dax.WorkerDiff{}, + expState: []dax.WorkerInfo{ + { + ID: "n2", + Jobs: []dax.Job{}, + }, + }, + }, + { + // Add a second worker. + fn: bal.AddWorker, + input: "n1", + expDiff: []dax.WorkerDiff{}, + expState: []dax.WorkerInfo{ + { + ID: "n1", + Jobs: []dax.Job{}, + }, + { + ID: "n2", + Jobs: []dax.Job{}, + }, + }, + }, + { + // Add job. + fn: bal.AddJob, + input: "p2", + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n1", + AddedJobs: []dax.Job{"p2"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n1", + Jobs: []dax.Job{"p2"}, + }, + { + ID: "n2", + Jobs: []dax.Job{}, + }, + }, + }, + { + // Add job. + fn: bal.AddJob, + input: "p3", + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n2", + AddedJobs: []dax.Job{"p3"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n1", + Jobs: []dax.Job{"p2"}, + }, + { + ID: "n2", + Jobs: []dax.Job{"p3"}, + }, + }, + }, + { + // Add job. + fn: bal.AddJob, + input: "p1", + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n1", + AddedJobs: []dax.Job{"p1"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n1", + Jobs: []dax.Job{"p1", "p2"}, + }, + { + ID: "n2", + Jobs: []dax.Job{"p3"}, + }, + }, + }, + { + // Add a third worker. + fn: bal.AddWorker, + input: "n0", + expDiff: []dax.WorkerDiff{}, + expState: []dax.WorkerInfo{ + { + ID: "n0", + Jobs: []dax.Job{}, + }, + { + ID: "n1", + Jobs: []dax.Job{"p1", "p2"}, + }, + { + ID: "n2", + Jobs: []dax.Job{"p3"}, + }, + }, + }, + { + // Add job. + fn: bal.AddJob, + input: "p4", + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n0", + AddedJobs: []dax.Job{"p4"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n0", + Jobs: []dax.Job{"p4"}, + }, + { + ID: "n1", + Jobs: []dax.Job{"p1", "p2"}, + }, + { + ID: "n2", + Jobs: []dax.Job{"p3"}, + }, + }, + }, + { + // Add job. + fn: bal.AddJob, + input: "p5", + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n0", + AddedJobs: []dax.Job{"p5"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n0", + Jobs: []dax.Job{"p4", "p5"}, + }, + { + ID: "n1", + Jobs: []dax.Job{"p1", "p2"}, + }, + { + ID: "n2", + Jobs: []dax.Job{"p3"}, + }, + }, + }, + { + // Add job. + fn: bal.AddJob, + input: "p0", + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n2", + AddedJobs: []dax.Job{"p0"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n0", + Jobs: []dax.Job{"p4", "p5"}, + }, + { + ID: "n1", + Jobs: []dax.Job{"p1", "p2"}, + }, + { + ID: "n2", + Jobs: []dax.Job{"p0", "p3"}, + }, + }, + }, + { + // Add job. + fn: bal.AddJob, + input: "p6", + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n0", + AddedJobs: []dax.Job{"p6"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n0", + Jobs: []dax.Job{"p4", "p5", "p6"}, + }, + { + ID: "n1", + Jobs: []dax.Job{"p1", "p2"}, + }, + { + ID: "n2", + Jobs: []dax.Job{"p0", "p3"}, + }, + }, + }, + { + // Add job. + fn: bal.AddJob, + input: "p7", + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n1", + AddedJobs: []dax.Job{"p7"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n0", + Jobs: []dax.Job{"p4", "p5", "p6"}, + }, + { + ID: "n1", + Jobs: []dax.Job{"p1", "p2", "p7"}, + }, + { + ID: "n2", + Jobs: []dax.Job{"p0", "p3"}, + }, + }, + }, + + //////////////////// Remove ///////////////////////// + + { + // Remove nonexistent worker. + fn: bal.RemoveWorker, + input: "nonexistent", + expDiff: []dax.WorkerDiff{}, + expState: []dax.WorkerInfo{ + { + ID: "n0", + Jobs: []dax.Job{"p4", "p5", "p6"}, + }, + { + ID: "n1", + Jobs: []dax.Job{"p1", "p2", "p7"}, + }, + { + ID: "n2", + Jobs: []dax.Job{"p0", "p3"}, + }, + }, + }, + { + // Remove worker. + fn: bal.RemoveWorker, + input: "n1", + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n1", + AddedJobs: []dax.Job{}, + RemovedJobs: []dax.Job{"p1", "p2", "p7"}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n0", + Jobs: []dax.Job{"p4", "p5", "p6"}, + }, + { + ID: "n2", + Jobs: []dax.Job{"p0", "p3"}, + }, + }, + }, + + { + // Remove job (from free list). + fn: bal.RemoveJob, + input: "p2", + expDiff: []dax.WorkerDiff{}, + expState: []dax.WorkerInfo{ + { + ID: "n0", + Jobs: []dax.Job{"p4", "p5", "p6"}, + }, + { + ID: "n2", + Jobs: []dax.Job{"p0", "p3"}, + }, + }, + }, + + { + // Balance after remove. + balance: true, + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n0", + AddedJobs: []dax.Job{"p7"}, + RemovedJobs: []dax.Job{}, + }, + { + WorkerID: "n2", + AddedJobs: []dax.Job{"p1"}, + RemovedJobs: []dax.Job{}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n0", + Jobs: []dax.Job{"p4", "p5", "p6", "p7"}, + }, + { + ID: "n2", + Jobs: []dax.Job{"p0", "p1", "p3"}, + }, + }, + }, + + { + // Remove job. + fn: bal.RemoveJob, + input: "p1", + expDiff: []dax.WorkerDiff{ + { + WorkerID: "n2", + AddedJobs: []dax.Job{}, + RemovedJobs: []dax.Job{"p1"}, + }, + }, + expState: []dax.WorkerInfo{ + { + ID: "n0", + Jobs: []dax.Job{"p4", "p5", "p6", "p7"}, + }, + { + ID: "n2", + Jobs: []dax.Job{"p0", "p3"}, + }, + }, + }, + } + for i, test := range tests { + t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) { + var diff []dax.WorkerDiff + var err error + if test.balance { + diff, err = bal.Balance(ctx) + } else { + diff, err = test.fn(ctx, newStringWrapper(test.input)) + } + assert.NoError(t, err) + assert.Equal(t, test.expDiff, diff) + + cs, err := bal.CurrentState(ctx) + assert.NoError(t, err) + assert.Equal(t, test.expState, cs) + }) + } + }) + + t.Run("WorkerState", func(t *testing.T) { + bal := boltdb.NewBalancer("test-worker-state", db, logger.NopLogger) + + _, err := bal.AddWorker(ctx, newStringWrapper("n1")) + assert.NoError(t, err) + _, err = bal.AddJob(ctx, newStringWrapper("p1")) + assert.NoError(t, err) + + exp := dax.WorkerInfo{ + ID: "n1", + Jobs: []dax.Job{"p1"}, + } + ws, err := bal.WorkerState(ctx, "n1") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + + // Worker doesn't exist. + exp = dax.WorkerInfo{ + ID: "x1", + } + ws, err = bal.WorkerState(ctx, "x1") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + }) + + t.Run("WorkersForJobs", func(t *testing.T) { + bal := boltdb.NewBalancer("test-workers-for-jobs", db, logger.NopLogger) + + _, err := bal.AddWorker(ctx, newStringWrapper("n1")) + assert.NoError(t, err) + _, err = bal.AddWorker(ctx, newStringWrapper("n2")) + assert.NoError(t, err) + for i := 0; i < 12; i++ { + _, err = bal.AddJob(ctx, newStringWrapper(fmt.Sprintf("p%d", i))) + assert.NoError(t, err) + } + + exp := dax.WorkerInfo{ + ID: "n1", + Jobs: []dax.Job{"p0", "p10", "p2", "p4", "p6", "p8"}, + } + ws, err := bal.WorkerState(ctx, "n1") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + + exp = dax.WorkerInfo{ + ID: "n2", + Jobs: []dax.Job{"p1", "p11", "p3", "p5", "p7", "p9"}, + } + ws, err = bal.WorkerState(ctx, "n2") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + + tests := []struct { + jobs []dax.Job + exp []dax.WorkerInfo + }{ + { + jobs: []dax.Job{"p0"}, + exp: []dax.WorkerInfo{ + {ID: "n1", Jobs: []dax.Job{"p0"}}, + }, + }, + { + jobs: []dax.Job{"p0", "p4"}, + exp: []dax.WorkerInfo{ + {ID: "n1", Jobs: []dax.Job{"p0", "p4"}}, + }, + }, + { + jobs: []dax.Job{"p0", "p4", "p999"}, + exp: []dax.WorkerInfo{ + {ID: "n1", Jobs: []dax.Job{"p0", "p4"}}, + }, + }, + { + jobs: []dax.Job{"p0", "p1"}, + exp: []dax.WorkerInfo{ + {ID: "n1", Jobs: []dax.Job{"p0"}}, + {ID: "n2", Jobs: []dax.Job{"p1"}}, + }, + }, + { + jobs: []dax.Job{"p5", "p0", "p1", "p8"}, + exp: []dax.WorkerInfo{ + {ID: "n1", Jobs: []dax.Job{"p0", "p8"}}, + {ID: "n2", Jobs: []dax.Job{"p1", "p5"}}, + }, + }, + } + for i, test := range tests { + t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) { + workers, err := bal.WorkersForJobs(ctx, test.jobs) + assert.NoError(t, err) + assert.Equal(t, test.exp, workers) + }) + } + }) + + t.Run("Balance", func(t *testing.T) { + bal := boltdb.NewBalancer("test-balance", db, logger.NopLogger) + + // Add two workers with some jobs evenly spread across them. + _, err := bal.AddWorker(ctx, newStringWrapper("n1")) + assert.NoError(t, err) + _, err = bal.AddWorker(ctx, newStringWrapper("n2")) + assert.NoError(t, err) + for i := 0; i < 13; i++ { + _, err = bal.AddJob(ctx, newStringWrapper(fmt.Sprintf("p%d", i))) + assert.NoError(t, err) + } + + exp := dax.WorkerInfo{ + ID: "n1", + Jobs: []dax.Job{"p0", "p10", "p12", "p2", "p4", "p6", "p8"}, + } + ws, err := bal.WorkerState(ctx, "n1") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + + exp = dax.WorkerInfo{ + ID: "n2", + Jobs: []dax.Job{"p1", "p11", "p3", "p5", "p7", "p9"}, + } + ws, err = bal.WorkerState(ctx, "n2") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + + // Now, add a worker and confirm that it currently has no jobs assigned + // to it. + _, err = bal.AddWorker(ctx, newStringWrapper("n3")) + assert.NoError(t, err) + exp = dax.WorkerInfo{ + ID: "n3", + Jobs: []dax.Job{}, + } + ws, err = bal.WorkerState(ctx, "n3") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + + // Finally, call Balance() and confirm that the appropriate jobs got + // reassigned. + _, err = bal.Balance(ctx) + assert.NoError(t, err) + + exp = dax.WorkerInfo{ + ID: "n1", + Jobs: []dax.Job{"p0", "p10", "p12", "p2", "p4"}, + } + ws, err = bal.WorkerState(ctx, "n1") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + + exp = dax.WorkerInfo{ + ID: "n2", + Jobs: []dax.Job{"p1", "p11", "p3", "p5"}, + } + ws, err = bal.WorkerState(ctx, "n2") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + + exp = dax.WorkerInfo{ + ID: "n3", + Jobs: []dax.Job{"p6", "p7", "p8", "p9"}, + } + ws, err = bal.WorkerState(ctx, "n3") + assert.NoError(t, err) + assert.Equal(t, exp, ws) + }) +} + +type stringWrapper struct { + s string +} + +func newStringWrapper(s string) *stringWrapper { + return &stringWrapper{ + s: s, + } +} + +func (s *stringWrapper) String() string { + return s.s +} diff --git a/dax/mds/controller/naive/types.go b/dax/mds/controller/naive/types.go new file mode 100644 index 000000000..f2679e781 --- /dev/null +++ b/dax/mds/controller/naive/types.go @@ -0,0 +1,81 @@ +package naive + +import ( + "sort" + + "github.com/molecula/featurebase/v3/dax" +) + +// jobSetDiffs is used internally to capture the diffs as they're happening. We +// call output() to generate the final result. +type jobSetDiffs struct { + added dax.Set[dax.Job] + removed dax.Set[dax.Job] +} + +func newJobSetDiffs() jobSetDiffs { + return jobSetDiffs{ + added: dax.NewSet[dax.Job](), + removed: dax.NewSet[dax.Job](), + } +} + +type internalDiffs map[dax.Worker]jobSetDiffs + +func newInternalDiffs() internalDiffs { + return make(internalDiffs) +} + +func (d internalDiffs) added(worker dax.Worker, job dax.Job) { + if _, ok := d[worker]; !ok { + d[worker] = newJobSetDiffs() + } + + // Before adding the job, make sure we haven't indicated that it has been + // removed prior to this. If it has, we need to invalidate that "remove" + // instruction. + d[worker].removed.Remove(job) + + d[worker].added.Add(job) +} + +func (d internalDiffs) removed(worker dax.Worker, job dax.Job) { + if _, ok := d[worker]; !ok { + d[worker] = newJobSetDiffs() + } + + // Before removing the job, make sure we haven't indicated that it has been + // added prior to this. If it has, we need to invalidate that "add" + // instruction. + d[worker].added.Remove(job) + + d[worker].removed.Add(job) +} + +func (d internalDiffs) merge(d2 internalDiffs) { + for k, v := range d2 { + if _, ok := d[k]; !ok { + d[k] = newJobSetDiffs() + } + d[k].added.Merge(v.added) + d[k].removed.Merge(v.removed) + } +} + +// output converts internalDiff to []controller.WorkerDiff for external +// consumption. +func (d internalDiffs) output() []dax.WorkerDiff { + out := make([]dax.WorkerDiff, len(d)) + + i := 0 + for k, v := range d { + out[i].WorkerID = k + out[i].AddedJobs = v.added.Sorted() + out[i].RemovedJobs = v.removed.Sorted() + i++ + } + + sort.Sort(dax.WorkerDiffs(out)) + + return out +} diff --git a/dax/mds/controller/node_registerer.go b/dax/mds/controller/node_registerer.go new file mode 100644 index 000000000..9c74c3ee3 --- /dev/null +++ b/dax/mds/controller/node_registerer.go @@ -0,0 +1,60 @@ +package controller + +import ( + "context" + "time" + + "github.com/molecula/featurebase/v3/dax" +) + +// nodeRegistrationRoutine is a long-running goroutine that reads +// newly registered nodes from a channel and sends out +// new directives to rebalance among all the nodes. +// +// If the provided timeout is 0, this routine will register each node placed on +// the channel immediately. +// +// If the provided timeout is >0, this routine will batch the nodes until the +// timeout time has passed. This prevents (for the case when scaling up by >1 +// node at a time) multiple directives being sent out serially as each node +// joins, and instead tries to handle all new nodes simultaneously. +func (c *Controller) nodeRegistrationRoutine(nodes chan *dax.Node, timeout time.Duration) error { + if timeout > 0 { + return c.nodeRegistrationDelayed(nodes, timeout) + } + return c.nodeRegistrationInstant(nodes) +} + +func (c *Controller) nodeRegistrationInstant(nodes chan *dax.Node) error { + for node := range nodes { + err := c.RegisterNodes(context.Background(), node) + if err != nil { + c.logger.Errorf("Registering node: %v, encountered error: %v", node, err) + } + } + return nil +} + +func (c *Controller) nodeRegistrationDelayed(nodes chan *dax.Node, timeout time.Duration) error { + batch := []*dax.Node{} + c.logger.Printf("Running with batch registration timeout: %v", timeout) + for { + select { + case <-c.stopping: + return nil + case node := <-nodes: + c.logger.Debugf("adding node: %+v", node) + batch = append(batch, node) + case <-time.After(timeout): + c.logger.Debugf("no new nodes in last %s, batch: %d", timeout, len(batch)) + if len(batch) > 0 { + err := c.RegisterNodes(context.Background(), batch...) + if err != nil { + c.logger.Errorf("Registering nodes: %v, encountered error: %v", batch, err) + } + + batch = batch[:0] // reset batch + } + } + } +} diff --git a/dax/mds/controller/partitioner/partitioner.go b/dax/mds/controller/partitioner/partitioner.go new file mode 100644 index 000000000..be35483a7 --- /dev/null +++ b/dax/mds/controller/partitioner/partitioner.go @@ -0,0 +1,55 @@ +// Package partitioner provides the Partitioner type, which provides helper +// methods for determining partitions based on string keys. +package partitioner + +import ( + "encoding/binary" + "hash/fnv" + + "github.com/molecula/featurebase/v3/dax" +) + +// Partitioner encapsulates helper methods for determining partitions +type Partitioner struct{} + +// NewPartitioner returns a new instance of Partitioner with default values. +func NewPartitioner() *Partitioner { + return &Partitioner{} +} + +// PartitionsForKeys returns a map of partitions to the list of strings which +// fall into that partition. +func (p *Partitioner) PartitionsForKeys(tkey dax.TableKey, partitionN int, keys ...string) map[dax.PartitionNum][]string { + out := make(map[dax.PartitionNum][]string) + + for _, key := range keys { + p := keyToPartition(tkey, partitionN, key) + if _, found := out[p]; !found { + out[p] = []string{} + } + out[p] = append(out[p], key) + } + + return out +} + +// keyToPartition returns the partition to which the given key belongs. +func keyToPartition(tkey dax.TableKey, partitionN int, key string) dax.PartitionNum { + // Hash the bytes and mod by partition count. + h := fnv.New64a() + _, _ = h.Write([]byte(tkey)) + _, _ = h.Write([]byte(key)) + return dax.PartitionNum(h.Sum64() % uint64(partitionN)) +} + +// ShardToPartition returns the PartitionNum for the given shard. +func (p *Partitioner) ShardToPartition(tkey dax.TableKey, shard dax.ShardNum, partitionN int) dax.PartitionNum { + var buf [8]byte + binary.BigEndian.PutUint64(buf[:], uint64(shard)) + + // Hash the bytes and mod by partition count. + h := fnv.New64a() + _, _ = h.Write([]byte(tkey)) + _, _ = h.Write(buf[:]) + return dax.PartitionNum(h.Sum64() % uint64(partitionN)) +} diff --git a/dax/mds/controller/partitioner/partitioner_test.go b/dax/mds/controller/partitioner/partitioner_test.go new file mode 100644 index 000000000..3e2fadf1f --- /dev/null +++ b/dax/mds/controller/partitioner/partitioner_test.go @@ -0,0 +1,86 @@ +package partitioner_test + +import ( + "fmt" + "testing" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/mds/controller/partitioner" + "github.com/stretchr/testify/assert" +) + +func TestPartitioner(t *testing.T) { + tableKey := dax.TableKey("foo") + partitionN := 8 + + t.Run("PartitionForKeys", func(t *testing.T) { + p := partitioner.NewPartitioner() + + tests := []struct { + tkey dax.TableKey + partitionN int + keys []string + exp map[dax.PartitionNum][]string + }{ + { + tkey: tableKey, + partitionN: partitionN, + keys: []string{"a"}, + exp: map[dax.PartitionNum][]string{ + 2: {"a"}, + }, + }, + { + tkey: tableKey, + partitionN: partitionN, + keys: []string{"a", "a"}, + exp: map[dax.PartitionNum][]string{ + 2: {"a", "a"}, + }, + }, + { + tkey: "differentTableName", + partitionN: partitionN, + keys: []string{"a"}, + exp: map[dax.PartitionNum][]string{ + 4: {"a"}, + }, + }, + { + tkey: tableKey, + partitionN: partitionN, + keys: []string{"a", "b", "c", "d", "e", "f", "g", "h", "i"}, + exp: map[dax.PartitionNum][]string{ + 0: {"g"}, + 1: {"d"}, + 2: {"a", "i"}, + 3: {"f"}, + 4: {"c"}, + 5: {"h"}, + 6: {"e"}, + 7: {"b"}, + }, + }, + } + for i, test := range tests { + t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) { + out := p.PartitionsForKeys(test.tkey, test.partitionN, test.keys...) + + assert.ElementsMatch(t, keys(test.exp), keys(out)) + + for k := range out { + assert.ElementsMatch(t, test.exp[k], out[k]) + } + }) + } + + }) +} + +func keys(m map[dax.PartitionNum][]string) dax.PartitionNums { + keys := make(dax.PartitionNums, 0, len(m)) + for k := range m { + keys = append(keys, k) + } + return keys +} diff --git a/dax/mds/controller/sets.go b/dax/mds/controller/sets.go new file mode 100644 index 000000000..d74856e69 --- /dev/null +++ b/dax/mds/controller/sets.go @@ -0,0 +1,169 @@ +package controller + +import ( + "sort" + + "github.com/molecula/featurebase/v3/dax" +) + +// StringSet is a set of strings. +type StringSet map[string]struct{} + +func NewStringSet() StringSet { + return make(StringSet) +} + +func (s StringSet) Add(p string) { + s[p] = struct{}{} +} + +func (s StringSet) Remove(p string) { + delete(s, p) +} + +func (s StringSet) Contains(p string) bool { + _, ok := s[p] + return ok +} + +func (s StringSet) SortedSlice() []string { + ps := make([]string, 0, len(s)) + for p := range s { + ps = append(ps, p) + } + sort.Strings(ps) + + return ps +} + +func (s StringSet) Minus(m StringSet) []string { + diff := []string{} + + for sk := range s { + var found bool + for mk := range m { + if mk == sk { + found = true + break + } + } + if !found { + diff = append(diff, sk) + } + } + + return diff +} + +// TableSet is a set of strings. +type TableSet map[dax.TableKey]struct{} + +func NewTableSet() TableSet { + return make(TableSet) +} + +func (s TableSet) Add(t dax.TableKey) { + s[t] = struct{}{} +} + +func (s TableSet) Remove(t dax.TableKey) { + delete(s, t) +} + +func (s TableSet) Contains(t dax.TableKey) bool { + _, ok := s[t] + return ok +} + +func (s TableSet) SortedSlice() dax.TableKeys { + ps := make(dax.TableKeys, 0, len(s)) + for p := range s { + ps = append(ps, p) + } + sort.Sort(ps) + + return ps +} + +func (s TableSet) QualifiedSortedSlice() map[dax.TableQualifier]dax.TableIDs { + m := make(map[dax.TableQualifier]dax.TableIDs) + for p := range s { + qtid := p.QualifiedTableID() + m[qtid.TableQualifier] = append(m[qtid.TableQualifier], qtid.ID) + } + + // Sort the slices in the map. + for _, v := range m { + sort.Sort(v) + } + + return m +} + +func (s TableSet) Minus(m TableSet) dax.TableKeys { + diff := dax.TableKeys{} + + for sk := range s { + var found bool + for mk := range m { + if mk == sk { + found = true + break + } + } + if !found { + diff = append(diff, sk) + } + } + + return diff +} + +// AddressSet is a set of strings. +type AddressSet map[dax.Address]struct{} + +func NewAddressSet() AddressSet { + return make(AddressSet) +} + +func (s AddressSet) Add(p dax.Address) { + s[p] = struct{}{} +} + +func (s AddressSet) Remove(p dax.Address) { + delete(s, p) +} + +func (s AddressSet) Contains(p dax.Address) bool { + _, ok := s[p] + return ok +} + +func (s AddressSet) SortedSlice() []dax.Address { + ps := make([]dax.Address, 0, len(s)) + for p := range s { + ps = append(ps, p) + } + sort.Slice(ps, func(i, j int) bool { return ps[i] < ps[j] }) + + return ps +} + +func (s AddressSet) Minus(m AddressSet) []dax.Address { + diff := []dax.Address{} + + for sk := range s { + var found bool + for mk := range m { + if mk == sk { + found = true + break + } + } + if !found { + diff = append(diff, sk) + } + } + + return diff +} diff --git a/dax/mds/controller/stringers.go b/dax/mds/controller/stringers.go new file mode 100644 index 000000000..30269f664 --- /dev/null +++ b/dax/mds/controller/stringers.go @@ -0,0 +1,104 @@ +package controller + +import ( + "fmt" + "strconv" + "strings" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/errors" +) + +// pUnit represents a table/partition combination. As a Stringer, it can be +// used as a job in the Balancer. +type pUnit struct { + t dax.TableKey + p dax.Partition +} + +func (p pUnit) String() string { + return fmt.Sprintf("%s|part_%d", p.t, p.p.Num) +} + +func (p pUnit) table() dax.TableKey { + return p.t +} + +func (p pUnit) partitionNum() dax.PartitionNum { + return p.p.Num +} + +func partition(t dax.TableKey, p dax.Partition) pUnit { + return pUnit{t, p} +} + +func decodePartition(j dax.Job) (pUnit, error) { + s := string(j) + parts := strings.Split(s, "|") + if len(parts) != 2 { + return pUnit{}, errors.Errorf("cannot decode string to partition: %s", s) + } + pparts := strings.Split(parts[1], "_") + if len(pparts) != 2 { + return pUnit{}, errors.Errorf("cannot decode partition part of string: %s", pparts[1]) + } + intVar, err := strconv.Atoi(pparts[1]) + if err != nil { + return pUnit{}, errors.Wrap(err, "converting string to int") + } + + return pUnit{ + t: dax.TableKey(parts[0]), + p: dax.Partition{ + Num: dax.PartitionNum(intVar), + Version: -1, + }, + }, nil +} + +// sUnit represents a table/shard combination. As a Stringer, it can be used as +// a job in the Balancer. +type sUnit struct { + t dax.TableKey + s dax.Shard +} + +func (s sUnit) String() string { + return fmt.Sprintf("%s|shard_%s", s.t, s.s.Num) +} + +func (s sUnit) table() dax.TableKey { + return s.t +} + +func (s sUnit) shardNum() dax.ShardNum { + return s.s.Num +} + +func shard(t dax.TableKey, s dax.Shard) sUnit { + return sUnit{t, s} +} + +func decodeShard(j dax.Job) (sUnit, error) { + s := string(j) + parts := strings.Split(s, "|") + if len(parts) != 2 { + return sUnit{}, errors.Errorf("cannot decode string to shardV: %s", s) + } + pparts := strings.Split(parts[1], "_") + if len(pparts) != 2 { + return sUnit{}, errors.Errorf("cannot decode shard part of string: %s", pparts[1]) + } + uint64Var, err := strconv.ParseUint(pparts[1], 10, 64) + if err != nil { + return sUnit{}, errors.Wrap(err, "converting string to int") + } + + return sUnit{ + t: dax.TableKey(parts[0]), + s: dax.Shard{ + Num: dax.ShardNum(uint64Var), + Version: -1, + }, + }, nil +} diff --git a/dax/mds/controller/types.go b/dax/mds/controller/types.go new file mode 100644 index 000000000..65d3e97a3 --- /dev/null +++ b/dax/mds/controller/types.go @@ -0,0 +1,19 @@ +package controller + +import "github.com/molecula/featurebase/v3/dax" + +// ComputeNode represents a compute node and the table/shards for which it is +// responsible. +type ComputeNode struct { + Address dax.Address `json:"address"` + Table dax.TableKey `json:"table"` + Shards dax.ShardNums `json:"shards"` +} + +// TranslateNode represents a translate node and the table/partitions for which +// it is responsible. +type TranslateNode struct { + Address dax.Address `json:"address"` + Table dax.TableKey `json:"table"` + Partitions dax.PartitionNums `json:"partitions"` +} diff --git a/dax/mds/http/addressmanager.go b/dax/mds/http/addressmanager.go new file mode 100644 index 000000000..0e9a770ca --- /dev/null +++ b/dax/mds/http/addressmanager.go @@ -0,0 +1,76 @@ +package http + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "log" + "net/http" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/errors" +) + +var ErrNotImplemented = errors.New(errors.ErrUncoded, "not implemented") + +// Ensure type implements interface. +var _ dax.AddressManager = &AddressManager{} + +// AddressManager is an http implementation of the AddressManager interface. +type AddressManager struct { + mdsAddress dax.Address +} + +func NewAddressManager(mdsAddress dax.Address) *AddressManager { + return &AddressManager{ + mdsAddress: mdsAddress, + } +} + +func (m *AddressManager) AddAddresses(ctx context.Context, addr ...dax.Address) error { + // Not implemented because it's currently not used + return ErrNotImplemented +} + +func (m *AddressManager) RemoveAddresses(ctx context.Context, addrs ...dax.Address) error { + if len(addrs) == 0 { + return nil + } + + if m.mdsAddress == "" { + return errors.Errorf("mdsAddress is empty; could not deregister: %s", addrs) + } + url := fmt.Sprintf("%s/deregister-nodes", m.mdsAddress.WithScheme("http")) + log.Printf("SEND deregister-nodes to: %s\n", url) + + req := DeregisterNodesRequest{ + Addresses: addrs, + } + + // Encode the request. + postBody, err := json.Marshal(req) + if err != nil { + return errors.Wrap(err, "marshalling deregister node request to json") + } + requestBody := bytes.NewBuffer(postBody) + + // Post the request. + request, _ := http.NewRequest(http.MethodPost, url, requestBody) + request.Header.Add("Content-Type", "application/json") + request.Header.Add("Accept", "application/json") + + resp, err := http.DefaultClient.Do(request) + if err != nil { + return errors.Wrap(err, "doing deregister node request") + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(resp.Body) + return errors.Errorf("status code: %d: %s", resp.StatusCode, b) + } + + return nil +} diff --git a/dax/mds/http/handler.go b/dax/mds/http/handler.go new file mode 100644 index 000000000..ff6fd3942 --- /dev/null +++ b/dax/mds/http/handler.go @@ -0,0 +1,664 @@ +package http + +import ( + "encoding/json" + "net/http" + + "github.com/gorilla/mux" + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/mds" + "github.com/molecula/featurebase/v3/dax/mds/controller" +) + +func Handler(mds *mds.MDS) http.Handler { + server := &server{ + mds: mds, + } + + router := mux.NewRouter() + router.HandleFunc("/health", server.getHealth).Methods("GET").Name("GetHealth") + + // mds endpoints. + router.HandleFunc("/create-table", server.postCreateTable).Methods("POST").Name("PostCreateTable") + router.HandleFunc("/drop-table", server.postDropTable).Methods("POST").Name("PostDropTable") + router.HandleFunc("/create-field", server.postCreateField).Methods("POST").Name("PostCreateField") + router.HandleFunc("/drop-field", server.postDropField).Methods("POST").Name("PostDropField") + router.HandleFunc("/table", server.postTable).Methods("POST").Name("PostTable") + router.HandleFunc("/table-id", server.postTableID).Methods("POST").Name("PostTable") + router.HandleFunc("/tables", server.postTables).Methods("POST").Name("PostTables") + + router.HandleFunc("/ingest-partition", server.postIngestPartition).Methods("POST").Name("PostIngestPartition") + router.HandleFunc("/ingest-shard", server.postIngestShard).Methods("POST").Name("PostIngestShard") + + router.HandleFunc("/snapshot", server.postSnapshot).Methods("POST").Name("PostSnapshot") + router.HandleFunc("/snapshot/shard-data", server.postSnapshotShardData).Methods("POST").Name("PostShapshotShardData") + router.HandleFunc("/snapshot/table-keys", server.postSnapshotTableKeys).Methods("POST").Name("PostShapshotTableKeys") + router.HandleFunc("/snapshot/field-keys", server.postSnapshotFieldKeys).Methods("POST").Name("PostShapshotFieldKeys") + + // controller endpoints. + router.HandleFunc("/register-node", server.postRegisterNode).Methods("POST").Name("PostRegisterNode") + router.HandleFunc("/register-nodes", server.postRegisterNodes).Methods("POST").Name("PostRegisterNodes") + router.HandleFunc("/deregister-nodes", server.postDeregisterNodes).Methods("POST").Name("PostDeregisterNodes") + router.HandleFunc("/check-in-node", server.postCheckInNode).Methods("POST").Name("PostCheckInNode") + router.HandleFunc("/compute-nodes", server.postComputeNodes).Methods("POST").Name("PostComputeNodes") + router.HandleFunc("/translate-nodes", server.postTranslateNodes).Methods("POST").Name("PostTranslateNodes") + + // debug endpoints + router.HandleFunc("/debug/nodes", server.getDebugNodes).Methods("GET").Name("GetDebugNodes") + + return router +} + +type server struct { + mds *mds.MDS +} + +// GET /health +func (s *server) getHealth(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) +} + +// POST /create-table +func (s *server) postCreateTable(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + ctx := r.Context() + + req := &dax.QualifiedTable{} + if err := json.NewDecoder(body).Decode(req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + err := s.mds.CreateTable(ctx, req) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + resp := CreateTableResponse(*req) + if err := json.NewEncoder(w).Encode(resp); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } +} + +type CreateTableResponse dax.QualifiedTable + +// POST /table +func (s *server) postTable(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + ctx := r.Context() + + qtid := dax.QualifiedTableID{} + if err := json.NewDecoder(body).Decode(&qtid); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + resp, err := s.mds.Table(ctx, qtid) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + if err := json.NewEncoder(w).Encode(resp); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } +} + +// POST /table-id +func (s *server) postTableID(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + ctx := r.Context() + + req := dax.QualifiedTableID{} + if err := json.NewDecoder(body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + qtid, err := s.mds.TableID(ctx, req.TableQualifier, req.Name) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + if err := json.NewEncoder(w).Encode(qtid); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } +} + +// POST /drop-table +func (s *server) postDropTable(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + ctx := r.Context() + + req := dax.QualifiedTableID{} + if err := json.NewDecoder(body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + err := s.mds.DropTable(ctx, req) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } +} + +// POST /create-field +func (s *server) postCreateField(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + ctx := r.Context() + + req := CreateFieldRequest{} + if err := json.NewDecoder(body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + qtid := req.TableKey.QualifiedTableID() + + err := s.mds.CreateField(ctx, qtid, req.Field) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } +} + +type CreateFieldRequest struct { + TableKey dax.TableKey `json:"table-key"` + Field *dax.Field `json:"field"` +} + +// POST /drop-field +func (s *server) postDropField(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + ctx := r.Context() + + req := DropFieldRequest{} + if err := json.NewDecoder(body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + qtid := req.Table + + err := s.mds.DropField(ctx, qtid, req.Field) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } +} + +type DropFieldRequest struct { + Table dax.QualifiedTableID `json:"table"` + Field dax.FieldName `json:"fields"` +} + +// POST /tables +func (s *server) postTables(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + ctx := r.Context() + + req := TablesRequest{} + if err := json.NewDecoder(body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + qual := dax.NewTableQualifier(req.OrganizationID, req.DatabaseID) + ids := req.TableIDs + + resp, err := s.mds.Tables(ctx, qual, ids...) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + if err := json.NewEncoder(w).Encode(resp); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } +} + +type TablesRequest struct { + OrganizationID dax.OrganizationID `json:"org-id"` + DatabaseID dax.DatabaseID `json:"db-id"` + TableIDs dax.TableIDs `json:"table-ids"` + TableNames dax.TableNames `json:"table-names"` +} + +// POST /ingest-partition +func (s *server) postIngestPartition(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + ctx := r.Context() + + req := IngestPartitionRequest{} + if err := json.NewDecoder(body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + qtid := req.Table + + addr, err := s.mds.IngestPartition(ctx, qtid, req.Partition) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + resp := &IngestPartitionResponse{ + Address: addr, + } + + if err := json.NewEncoder(w).Encode(resp); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } +} + +type IngestPartitionRequest struct { + Table dax.QualifiedTableID `json:"table"` + Partition dax.PartitionNum `json:"partition"` +} + +type IngestPartitionResponse struct { + Address dax.Address `json:"address"` +} + +// POST /ingest-shard +func (s *server) postIngestShard(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + ctx := r.Context() + + req := IngestShardRequest{} + if err := json.NewDecoder(body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + qtid := req.Table + + addr, err := s.mds.IngestShard(ctx, qtid, req.Shard) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + resp := &IngestShardResponse{ + Address: addr, + } + + if err := json.NewEncoder(w).Encode(resp); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } +} + +type IngestShardRequest struct { + Table dax.QualifiedTableID `json:"table"` + Shard dax.ShardNum `json:"shard"` +} + +type IngestShardResponse struct { + Address dax.Address `json:"address"` +} + +// POST /snapshot +// High level snapshot endpoint to snapshot everything in a table. +func (s *server) postSnapshot(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + ctx := r.Context() + + req := dax.QualifiedTableID{} + if err := json.NewDecoder(body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + if err := s.mds.SnapshotTable(ctx, req); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + w.WriteHeader(http.StatusOK) +} + +// POST /snapshot/shard-data +func (s *server) postSnapshotShardData(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + ctx := r.Context() + + req := SnapshotShardRequest{} + if err := json.NewDecoder(body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + qtid := req.Table + + if err := s.mds.SnapshotShardData(ctx, qtid, req.Shard); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + w.WriteHeader(http.StatusOK) +} + +// SnapshotShardRequest is used to specify the table/shard to snapshot. +type SnapshotShardRequest struct { + Table dax.QualifiedTableID `json:"table"` + Shard dax.ShardNum `json:"shard"` +} + +// POST /snapshot/table-keys +func (s *server) postSnapshotTableKeys(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + ctx := r.Context() + + req := SnapshotTableKeysRequest{} + if err := json.NewDecoder(body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + qtid := req.Table + + if err := s.mds.SnapshotTableKeys(ctx, qtid, req.Partition); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + w.WriteHeader(http.StatusOK) +} + +// SnapshotTableKeysRequest is used to specify the table/partition/keys to +// snapshot. +type SnapshotTableKeysRequest struct { + Table dax.QualifiedTableID `json:"table"` + Partition dax.PartitionNum `json:"partition"` +} + +// POST /snapshot/field-keys +func (s *server) postSnapshotFieldKeys(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + ctx := r.Context() + + req := SnapshotFieldKeysRequest{} + if err := json.NewDecoder(body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + qtid := req.Table + + if err := s.mds.SnapshotFieldKeys(ctx, qtid, req.Field); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + w.WriteHeader(http.StatusOK) +} + +// SnapshotFieldKeysRequest is used to specify the field/keys to snapshot. +type SnapshotFieldKeysRequest struct { + Table dax.QualifiedTableID `json:"table"` + Field dax.FieldName `json:"field"` +} + +// POST /register-node +func (s *server) postRegisterNode(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + ctx := r.Context() + + req := RegisterNodeRequest{} + if err := json.NewDecoder(body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + node := &dax.Node{ + Address: req.Address, + RoleTypes: req.RoleTypes, + } + + if err := s.mds.RegisterNode(ctx, node); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + w.WriteHeader(http.StatusOK) +} + +type RegisterNodeRequest struct { + Address dax.Address `json:"address"` + + // RoleTypes allows a registering node to specify which role type(s) it is + // capable of filling. The controller will not assign a role to this node + // with a type not included in RoleTypes. + RoleTypes []dax.RoleType `json:"role-types"` +} + +// POST /register-nodes +func (s *server) postRegisterNodes(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + ctx := r.Context() + + req := RegisterNodesRequest{} + if err := json.NewDecoder(body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + if err := s.mds.RegisterNodes(ctx, req.Nodes...); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + w.WriteHeader(http.StatusOK) +} + +type RegisterNodesRequest struct { + Nodes []*dax.Node `json:"nodes"` +} + +// POST /deregister-nodes +func (s *server) postDeregisterNodes(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + ctx := r.Context() + + req := DeregisterNodesRequest{} + if err := json.NewDecoder(body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + if err := s.mds.DeregisterNodes(ctx, req.Addresses...); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + w.WriteHeader(http.StatusOK) +} + +type DeregisterNodesRequest struct { + Addresses []dax.Address `json:"addresses"` +} + +// POST /check-in-node +func (s *server) postCheckInNode(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + ctx := r.Context() + + req := CheckInNodeRequest{} + if err := json.NewDecoder(body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + node := &dax.Node{ + Address: req.Address, + RoleTypes: req.RoleTypes, + } + + if err := s.mds.CheckInNode(ctx, node); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + w.WriteHeader(http.StatusOK) +} + +type CheckInNodeRequest struct { + Address dax.Address `json:"address"` + + // RoleTypes allows a registering node to specify which role type(s) it is + // capable of filling. The controller will not assign a role to this node + // with a type not included in RoleTypes. + RoleTypes []dax.RoleType `json:"role-types"` +} + +// POST /compute-nodes +func (s *server) postComputeNodes(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + ctx := r.Context() + + req := ComputeNodesRequest{} + if err := json.NewDecoder(body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + qtid := req.Table + + nodes, err := s.mds.ComputeNodes(ctx, qtid, req.Shards...) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + resp := ComputeNodesResponse{ + ComputeNodes: nodes, + } + + if err := json.NewEncoder(w).Encode(resp); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } +} + +func (s *server) getDebugNodes(w http.ResponseWriter, r *http.Request) { + nodes, err := s.mds.DebugNodes(r.Context()) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + if err := json.NewEncoder(w).Encode(nodes); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + +} + +// ComputeNodesRequest is used to specify the table/shards to consider in the +// ComputeNodes method call. If IsWrite is true, shards which are not currently +// being managed by the underlying Controller will be added to (registered with) +// the Controller and, if adequate compute is available, will be associated with +// a compute node. +type ComputeNodesRequest struct { + Table dax.QualifiedTableID `json:"table"` + Shards dax.ShardNums `json:"shards"` + IsWrite bool `json:"is-write"` +} + +// ComputeNodesResponse contains the list of compute nodes returned based on the +// table/shards specified in the ComputeNodeRequest. It's possible that shards +// provided are not included in this response. That might happen if there are +// currently no active compute nodes. +type ComputeNodesResponse struct { + ComputeNodes []controller.ComputeNode `json:"compute-nodes"` +} + +// POST /translate-nodes +func (s *server) postTranslateNodes(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + ctx := r.Context() + + req := TranslateNodesRequest{} + if err := json.NewDecoder(body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + qtid := req.Table + + nodes, err := s.mds.TranslateNodes(ctx, qtid, req.Partitions...) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + resp := TranslateNodesResponse{ + TranslateNodes: nodes, + } + + if err := json.NewEncoder(w).Encode(resp); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } +} + +// TranslateNodesRequest is used to specify the table/partitions to consider in +// the TranslateNodes method call. +type TranslateNodesRequest struct { + Table dax.QualifiedTableID `json:"table"` + Partitions dax.PartitionNums `json:"partitions"` + IsWrite bool `json:"is-write"` +} + +// TranslateNodesResponse contains the list of translate nodes returned based on +// the table/partitions specified in the TranslateNodeRequest. It's possible +// that partitions provided are not included in this response. That might happen +// if there are currently no active translate nodes. +type TranslateNodesResponse struct { + TranslateNodes []controller.TranslateNode `json:"translate-nodes"` +} diff --git a/dax/mds/mds.go b/dax/mds/mds.go new file mode 100644 index 000000000..04a331115 --- /dev/null +++ b/dax/mds/mds.go @@ -0,0 +1,463 @@ +// Package mds provides the overall interface to Metadata Services. +package mds + +import ( + "context" + "fmt" + "log" + "os" + "sync" + "time" + + fb "github.com/molecula/featurebase/v3" + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/boltdb" + "github.com/molecula/featurebase/v3/dax/mds/controller" + naiveboltdb "github.com/molecula/featurebase/v3/dax/mds/controller/naive/boltdb" + "github.com/molecula/featurebase/v3/dax/mds/poller" + "github.com/molecula/featurebase/v3/dax/mds/schemar" + schemarboltdb "github.com/molecula/featurebase/v3/dax/mds/schemar/boltdb" + "github.com/molecula/featurebase/v3/errors" + "github.com/molecula/featurebase/v3/logger" +) + +type Config struct { + // Controller + Director controller.Director + // RegistrationBatchTimeout is the time that the controller will + // wait after a node registers itself to see if any more nodes + // will register before sending out directives to all nodes which + // have been registered. + RegistrationBatchTimeout time.Duration + + // Poller + PollInterval time.Duration + + // Storage + StorageMethod string + StorageDSN string + + // Logger + Logger logger.Logger +} + +// Ensure type implements interface. +var _ fb.MDS = (*MDS)(nil) + +// MDS provides public MDS methods for an MDS service. +type MDS struct { + mu sync.RWMutex + + controller *controller.Controller + poller *poller.Poller + schemar schemar.Schemar + + logger logger.Logger +} + +// New returns a new instance of MDS. +func New(cfg Config) *MDS { + // Set up logger. + var logr = logger.NopLogger + if cfg.Logger != nil { + logr = cfg.Logger + } + + // Storage methods. + if cfg.StorageMethod != "boltdb" && cfg.StorageMethod != "" { + log.Printf("storagemethod %s not supported, try 'boltdb'", cfg.StorageMethod) + } + if cfg.StorageDSN == "" { + dir, err := os.MkdirTemp("", "mds_*") + if err != nil { + logr.Printf("Making temp dir for MDS storage: %v", err) + os.Exit(1) + } + cfg.StorageDSN = fmt.Sprintf("file:%s", dir) + logr.Warnf("no StorageDSN given (like 'file:/path/to/directory') using temp dir at '%s'", cfg.StorageDSN) + } + schemarDB, err := boltdb.NewSvcBolt(cfg.StorageDSN, "schemar", schemarboltdb.SchemarBuckets...) + if err != nil { + logr.Printf("Error creating schemar db: %v", err) + os.Exit(1) + } + schemar := schemarboltdb.NewSchemar(schemarDB, logr) + + boltDB, err := boltdb.NewSvcBolt(cfg.StorageDSN, "balancer", naiveboltdb.NaiveBalancerBuckets...) + if err != nil { + log.Println(errors.Wrap(err, "creating balancer bolt")) + os.Exit(1) + } + + controllerCfg := controller.Config{ + Director: cfg.Director, + Schemar: schemar, + ComputeBalancer: naiveboltdb.NewBalancer("compute", boltDB, logr), + TranslateBalancer: naiveboltdb.NewBalancer("translate", boltDB, logr), + + RegistrationBatchTimeout: cfg.RegistrationBatchTimeout, + + StorageMethod: cfg.StorageMethod, + // just reusing this bolt for internal controller svcs + // rn... ultimately controller shouldn't know what bolt is at + // all + BoltDB: boltDB, + + Logger: logr, + } + controller := controller.New(controllerCfg) + + pollerCfg := poller.Config{ + AddressManager: controller, + NodePoller: poller.NewHTTPNodePoller(logr), + PollInterval: cfg.PollInterval, + Logger: logr, + } + poller := poller.New(pollerCfg) + + // The controller needs to tell the poller about nodes which have been + // added/removed. + // TODO: this feels hacky. We need an elegant way to register interface + // implementations across services without an explicit Set method like this. + controller.SetPoller(poller) + + return &MDS{ + controller: controller, + poller: poller, + schemar: schemar, + + logger: logr, + } +} + +//////////////////////////////////////////////////// +// mds specific endpoints +//////////////////////////////////////////////////// + +// Run starts MDS services, such as the Poller. +func (m *MDS) Run() error { + // Initialize the poller (in the case where this MDS instance has restarted + // or is a replacement). Then start the poller. + if err := m.controller.InitializePoller(context.Background()); err != nil { + return errors.Wrap(err, "initializing the poller") + } + m.poller.Run() + + return m.controller.Run() +} + +// Stop stops MDS services, such as the Poller and the controller's node +// registration routine. +func (m *MDS) Stop() error { + m.poller.Stop() + m.controller.Stop() + return nil +} + +// sanitizeQTID populates Table.ID (by looking up the table, by name, in +// schemar) for a given table having only a Name value, but no ID. +func (m *MDS) sanitizeQTID(ctx context.Context, qtid *dax.QualifiedTableID) error { + if qtid.ID == "" { + nqtid, err := m.schemar.TableID(ctx, qtid.TableQualifier, qtid.Name) + if err != nil { + return errors.Wrap(err, "getting table ID") + } + qtid.ID = nqtid.ID + } + return nil +} + +// CreateTable handles a create table request. +func (m *MDS) CreateTable(ctx context.Context, qtbl *dax.QualifiedTable) error { + m.mu.Lock() + defer m.mu.Unlock() + + // Create Table ID. + if _, err := qtbl.CreateID(); err != nil { + return errors.Wrap(err, "creating table ID") + } + + // Create the table in schemar. + if err := m.schemar.CreateTable(ctx, qtbl); err != nil { + return errors.Wrapf(err, "creating table: %s", qtbl) + } + + // TODO: if error here, we should probably roll-back the + // schemar.CreateTable() request. + + // Add the table to the controller. + return m.controller.CreateTable(ctx, qtbl) +} + +// DropTable handles a drop table request. // TODO(jaffee) how do we +// reason about consistency here? What if controller DropTable +// succeeds, but schemar fails? +func (m *MDS) DropTable(ctx context.Context, qtid dax.QualifiedTableID) error { + m.mu.Lock() + defer m.mu.Unlock() + + if err := m.sanitizeQTID(ctx, &qtid); err != nil { + return errors.Wrap(err, "sanitizing") + } + + if err := m.controller.DropTable(ctx, qtid); err != nil { + return errors.Wrapf(err, "dropping table: %s", qtid) + } + + return m.schemar.DropTable(ctx, qtid) +} + +type CreateFieldRequest struct { + Table dax.TableName + Field *dax.Field +} + +// CreateField handles a create Field request. +func (m *MDS) CreateField(ctx context.Context, qtid dax.QualifiedTableID, fld *dax.Field) error { + m.mu.Lock() + defer m.mu.Unlock() + + if err := m.sanitizeQTID(ctx, &qtid); err != nil { + return errors.Wrap(err, "sanitizing") + } + + // Create the field in schemar. + if err := m.schemar.CreateField(ctx, qtid, fld); err != nil { + return errors.Wrapf(err, "creating field: %s, %s", qtid, fld) + } + + // Add the table to the controller. + return m.controller.CreateField(ctx, qtid, fld) +} + +// DropField handles a drop Field request. +func (m *MDS) DropField(ctx context.Context, qtid dax.QualifiedTableID, fldName dax.FieldName) error { + m.mu.Lock() + defer m.mu.Unlock() + + if err := m.sanitizeQTID(ctx, &qtid); err != nil { + return errors.Wrap(err, "sanitizing") + } + + // Drop the field from schemar. + if err := m.schemar.DropField(ctx, qtid, fldName); err != nil { + return errors.Wrapf(err, "dropping field: %s, %s", qtid, fldName) + } + + // Drop the field from the controller. + return m.controller.DropField(ctx, qtid, fldName) +} + +type DropFieldResponse struct{} + +// Table handles a table request. +func (m *MDS) Table(ctx context.Context, qtid dax.QualifiedTableID) (*dax.QualifiedTable, error) { + m.mu.RLock() + defer m.mu.RUnlock() + + if err := m.sanitizeQTID(ctx, &qtid); err != nil { + return nil, errors.Wrap(err, "sanitizing") + } + + return m.schemar.Table(ctx, qtid) +} + +// Tables handles a tables request. +func (m *MDS) Tables(ctx context.Context, qual dax.TableQualifier, ids ...dax.TableID) ([]*dax.QualifiedTable, error) { + m.mu.RLock() + defer m.mu.RUnlock() + + return m.schemar.Tables(ctx, qual, ids...) +} + +// TableID handles a table id (i.e. by name) request. +func (m *MDS) TableID(ctx context.Context, qual dax.TableQualifier, name dax.TableName) (dax.QualifiedTableID, error) { + m.mu.RLock() + defer m.mu.RUnlock() + + return m.schemar.TableID(ctx, qual, name) +} + +// IngestPartition handles an ingest partition request. +func (m *MDS) IngestPartition(ctx context.Context, qtid dax.QualifiedTableID, partnNum dax.PartitionNum) (dax.Address, error) { + m.mu.RLock() + defer m.mu.RUnlock() + + if err := m.sanitizeQTID(ctx, &qtid); err != nil { + return "", errors.Wrap(err, "sanitizing") + } + + // Verify that the table exists. + if _, err := m.schemar.Table(ctx, qtid); err != nil { + return "", err + } + + partitions := dax.PartitionNums{partnNum} + + nodes, err := m.controller.TranslateNodes(ctx, qtid, partitions, true) + if err != nil { + return "", err + } + + if l := len(nodes); l == 0 { + return "", controller.NewErrNoAvailableNode() + } else if l > 1 { + return "", controller.NewErrInternal( + fmt.Sprintf("unexpected number of nodes: %d", l)) + } + + node := nodes[0] + + // Verify that the node returned is actually responsible for the partition + // requested. + if node.Table != qtid.Key() { + return "", controller.NewErrInternal( + fmt.Sprintf("table returned (%s) does not match requested (%s)", node.Table, qtid)) + } else if l := len(node.Partitions); l != 1 { + return "", controller.NewErrInternal( + fmt.Sprintf("unexpected number of partitions returned: %d", l)) + } else if p := node.Partitions[0]; p != partnNum { + return "", controller.NewErrInternal( + fmt.Sprintf("partition returned (%d) does not match requested (%d)", p, partnNum)) + } + + return node.Address, nil +} + +// IngestShard handles an ingest shard request. +func (m *MDS) IngestShard(ctx context.Context, qtid dax.QualifiedTableID, shrdNum dax.ShardNum) (dax.Address, error) { + m.mu.RLock() + defer m.mu.RUnlock() + + if err := m.sanitizeQTID(ctx, &qtid); err != nil { + return "", errors.Wrap(err, "sanitizing") + } + + // Verify that the table exists. + if _, err := m.schemar.Table(ctx, qtid); err != nil { + return "", err + } + + shards := dax.ShardNums{shrdNum} + + nodes, err := m.controller.ComputeNodes(ctx, qtid, shards, true) + if err != nil { + return "", err + } + + if l := len(nodes); l == 0 { + return "", controller.NewErrNoAvailableNode() + } else if l > 1 { + return "", controller.NewErrInternal( + fmt.Sprintf("unexpected number of nodes: %d", l)) + } + + node := nodes[0] + + // Verify that the node returned is actually responsible for the shard + // requested. + if node.Table != qtid.Key() { + return "", controller.NewErrInternal( + fmt.Sprintf("table returned (%s) does not match requested (%s)", node.Table, qtid)) + } else if l := len(node.Shards); l != 1 { + return "", controller.NewErrInternal( + fmt.Sprintf("unexpected number of shards returned: %d", l)) + } else if s := node.Shards[0]; s != shrdNum { + return "", controller.NewErrInternal( + fmt.Sprintf("shard returned (%d) does not match requested (%d)", s, shrdNum)) + } + + return node.Address, nil +} + +// SnapshotTable handles a snapshot table request. +func (m *MDS) SnapshotTable(ctx context.Context, qtid dax.QualifiedTableID) error { + if err := m.sanitizeQTID(ctx, &qtid); err != nil { + return errors.Wrap(err, "sanitizing") + } + + return m.controller.SnapshotTable(ctx, qtid) +} + +// SnapshotShardData handles a snapshot shard request. +func (m *MDS) SnapshotShardData(ctx context.Context, qtid dax.QualifiedTableID, shardNum dax.ShardNum) error { + if err := m.sanitizeQTID(ctx, &qtid); err != nil { + return errors.Wrap(err, "sanitizing") + } + + return m.controller.SnapshotShardData(ctx, qtid, shardNum) +} + +// SnapshotTableKeys handles a snapshot table/keys request. +func (m *MDS) SnapshotTableKeys(ctx context.Context, qtid dax.QualifiedTableID, partitionNum dax.PartitionNum) error { + if err := m.sanitizeQTID(ctx, &qtid); err != nil { + return errors.Wrap(err, "sanitizing") + } + + return m.controller.SnapshotTableKeys(ctx, qtid, partitionNum) +} + +// SnapshotFieldKeys handles a snapshot field/keys request. +func (m *MDS) SnapshotFieldKeys(ctx context.Context, qtid dax.QualifiedTableID, fldName dax.FieldName) error { + if err := m.sanitizeQTID(ctx, &qtid); err != nil { + return errors.Wrap(err, "sanitizing") + } + + return m.controller.SnapshotFieldKeys(ctx, qtid, fldName) +} + +//////////////////////////////////////////////////// +// controller specific endpoints +// These are just pass-throughs for now. +//////////////////////////////////////////////////// + +// RegisterNode handles a node registration request. It does not +// synchronously do much of anything, but the node will eventually +// probably get a directive... unless the MDS crashes or something in +// which case the fact that this endpoint was ever called will be lost +// to time. +func (m *MDS) RegisterNode(ctx context.Context, node *dax.Node) error { + return m.controller.RegisterNode(ctx, node) +} + +// CheckInNode handles a node check-in request. If MDS is not aware of the node, +// it will be sent through the RegisterNode process. +func (m *MDS) CheckInNode(ctx context.Context, node *dax.Node) error { + return m.controller.CheckInNode(ctx, node) +} + +// RegisterNodes immediately registers the given nodes and sends out +// new directives synchronously, bypassing the wait time of the +// RegisterNode endpoint. +func (m *MDS) RegisterNodes(ctx context.Context, nodes ...*dax.Node) error { + return m.controller.RegisterNodes(ctx, nodes...) +} + +// DeregisterNodes handles a request to deregister multiple nodes at once. +func (m *MDS) DeregisterNodes(ctx context.Context, addrs ...dax.Address) error { + return m.controller.DeregisterNodes(ctx, addrs...) +} + +// ComputeNodes gets the compute nodes responsible for the table/shards +// specified in the ComputeNodeRequest. +func (m *MDS) ComputeNodes(ctx context.Context, qtid dax.QualifiedTableID, shardNums ...dax.ShardNum) ([]controller.ComputeNode, error) { + if err := m.sanitizeQTID(ctx, &qtid); err != nil { + return nil, errors.Wrap(err, "sanitizing") + } + + return m.controller.ComputeNodes(ctx, qtid, shardNums, false) +} + +func (m *MDS) DebugNodes(ctx context.Context) ([]*dax.Node, error) { + return m.controller.DebugNodes(ctx) +} + +// TranslateNodes gets the translate nodes responsible for the table/partitions +// specified in the TranslateNodeRequest. +func (m *MDS) TranslateNodes(ctx context.Context, qtid dax.QualifiedTableID, partitionNums ...dax.PartitionNum) ([]controller.TranslateNode, error) { + if err := m.sanitizeQTID(ctx, &qtid); err != nil { + return nil, errors.Wrap(err, "sanitizing") + } + + return m.controller.TranslateNodes(ctx, qtid, partitionNums, false) +} diff --git a/dax/mds/poller/config.go b/dax/mds/poller/config.go new file mode 100644 index 000000000..b259008a7 --- /dev/null +++ b/dax/mds/poller/config.go @@ -0,0 +1,15 @@ +package poller + +import ( + "time" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/logger" +) + +type Config struct { + AddressManager dax.AddressManager + NodePoller NodePoller + PollInterval time.Duration + Logger logger.Logger +} diff --git a/dax/mds/poller/interfaces.go b/dax/mds/poller/interfaces.go new file mode 100644 index 000000000..2eabaf045 --- /dev/null +++ b/dax/mds/poller/interfaces.go @@ -0,0 +1,59 @@ +package poller + +import ( + "fmt" + "net/http" + "time" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/logger" +) + +// NodePoller is an interface to anything which has the ability to poll a +// node. +type NodePoller interface { + Poll(dax.Address) bool +} + +// Ensure type implements interface. +var _ NodePoller = (*NopNodePoller)(nil) +var _ NodePoller = (*HTTPNodePoller)(nil) + +// NopNodePoller is a no-op implementation of the NodePoller interface. +type NopNodePoller struct{} + +func NewNopNodePoller() *NopNodePoller { + return &NopNodePoller{} +} + +func (p *NopNodePoller) Poll(addr dax.Address) bool { + return true +} + +// HTTPNodePoller is an http implementation of the NodePoller interface. +type HTTPNodePoller struct { + logger logger.Logger + client *http.Client +} + +func NewHTTPNodePoller(logger logger.Logger) *HTTPNodePoller { + return &HTTPNodePoller{ + logger: logger, + client: &http.Client{ + Timeout: time.Second, // short timeout for polling to detect issues quickly. /health endpoints should always respond fast. + }, + } +} + +func (p *HTTPNodePoller) Poll(addr dax.Address) bool { + url := fmt.Sprintf("%s/health", addr.WithScheme("http")) + + if resp, err := p.client.Get(url); err != nil { + p.logger.Printf("poll error: %s\n", err) + return false + } else if resp.StatusCode != http.StatusOK { + return false + } + + return true +} diff --git a/dax/mds/poller/poller.go b/dax/mds/poller/poller.go new file mode 100644 index 000000000..51d5f28ed --- /dev/null +++ b/dax/mds/poller/poller.go @@ -0,0 +1,156 @@ +// Package poller provides the core Poller struct. +package poller + +import ( + "context" + "sync" + "time" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/logger" +) + +// Poller maintains a list of nodes to poll. It also polls them. +type Poller struct { + mu sync.RWMutex + + addresses map[dax.Address]struct{} + + addressManager dax.AddressManager + + nodePoller NodePoller + pollInterval time.Duration + + running bool + stopping chan struct{} + + logger logger.Logger +} + +// New returns a new instance of Poller with default values. +func New(cfg Config) *Poller { + p := &Poller{ + addresses: make(map[dax.Address]struct{}), + addressManager: dax.NewNopAddressManager(), + nodePoller: NewNopNodePoller(), + pollInterval: time.Second, + stopping: make(chan struct{}), + logger: logger.NopLogger, + } + + // Set config options. + if cfg.AddressManager != nil { + p.addressManager = cfg.AddressManager + } + if cfg.NodePoller != nil { + p.nodePoller = cfg.NodePoller + } + if cfg.PollInterval != 0 { + p.pollInterval = cfg.PollInterval + } + if cfg.Logger != nil { + p.logger = cfg.Logger + } + + return p +} + +func (p *Poller) AddAddresses(ctx context.Context, addrs ...dax.Address) error { + p.mu.Lock() + defer p.mu.Unlock() + + for _, addr := range addrs { + p.addresses[addr] = struct{}{} + } + + return nil +} + +func (p *Poller) RemoveAddresses(ctx context.Context, addrs ...dax.Address) error { + p.mu.Lock() + defer p.mu.Unlock() + + for _, addr := range addrs { + delete(p.addresses, addr) + } + + return nil +} + +func (p *Poller) Addresses() []dax.Address { + p.mu.RLock() + defer p.mu.RUnlock() + + addrs := make([]dax.Address, 0, len(p.addresses)) + for addr := range p.addresses { + addrs = append(addrs, addr) + } + + return addrs +} + +// Run starts the polling goroutine. +func (p *Poller) Run() { + p.mu.Lock() + defer p.mu.Unlock() + + if p.running { + p.logger.Printf("poller is already running") + return + } + p.running = true + + go func() { p.run() }() +} + +func (p *Poller) run() { + ticker := time.NewTicker(p.pollInterval) + defer ticker.Stop() + + for { + // Wait for tick or a close. + select { + case <-p.stopping: + return + case <-ticker.C: + } + + p.pollAll() + } + +} + +// Stop stops the polling routine. +func (p *Poller) Stop() { + close(p.stopping) +} + +func (p *Poller) pollAll() { + addrs := p.Addresses() + + ctx := context.Background() + + toRemove := []dax.Address{} + + for _, addr := range addrs { + p.logger.Debugf("polling: %s", addr) + start := time.Now() + up := p.nodePoller.Poll(addr) + if !up { + p.logger.Printf("poller removing %s", addr) + toRemove = append(toRemove, addr) + } + p.logger.Debugf("done poll: %s, %s", addr, time.Since(start)) + } + + if len(toRemove) > 0 { + p.logger.Debugf("removing addresses: %v", toRemove) + start := time.Now() + err := p.addressManager.RemoveAddresses(ctx, toRemove...) + if err != nil { + p.logger.Printf("removing %s: %v", toRemove, err) + } + p.logger.Debugf("remove complete: %s", time.Since(start)) + } + +} diff --git a/dax/mds/poller/poller_test.go b/dax/mds/poller/poller_test.go new file mode 100644 index 000000000..ceb5e2470 --- /dev/null +++ b/dax/mds/poller/poller_test.go @@ -0,0 +1,182 @@ +package poller_test + +import ( + "context" + "encoding/json" + "log" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/molecula/featurebase/v3/dax" + mds_http "github.com/molecula/featurebase/v3/dax/mds/http" + "github.com/molecula/featurebase/v3/dax/mds/poller" + "github.com/molecula/featurebase/v3/logger" + "github.com/stretchr/testify/assert" +) + +// TestPoller runs for a total of 5 seconds. It begins by polling two healthy +// nodes. After 3 seconds, one of the nodes dies. At that point, the poller +// de-registers the dead node from the node manager, after which the node +// manager tells the poller to stop polling the dead node. +func TestPoller(t *testing.T) { + ctx := context.Background() + + // node 1 + node1 := newMockNode(t, "health", 0) + defer node1.Close() + addr1 := dax.Address(node1.URL()) + + // node 1 + node2 := newMockNode(t, "health", 3*time.Second) + defer node2.Close() + addr2 := dax.Address(node2.URL()) + + // manager + manager := newMockManager(t, ctx, "deregister-nodes", []dax.Address{addr1, addr2}) + defer manager.Close() + managerAddr := dax.Address(manager.URL()) + + t.Run("Poller", func(t *testing.T) { + cfg := poller.Config{ + AddressManager: mds_http.NewAddressManager(managerAddr), + NodePoller: poller.NewHTTPNodePoller(logger.NopLogger), + } + p := poller.New(cfg) + + // This is a little strange, but basically we need the manager to be + // able to call poller.RemoveAddresses, and since this test poller isn't + // running as an http server (unlike everything else in this test: + // manager, nodes), we give the manager a pointer to the Poller here so + // it can call the RemoveAddresses method directly. + manager.setPoller(p) + + done := make(chan struct{}) + go func() { + time.Sleep(5 * time.Second) + close(done) + }() + + p.Run() + defer p.Stop() + + p.AddAddresses(ctx, addr1, addr2) + + // wait for a done + <-done + + assert.Contains(t, p.Addresses(), addr1) + assert.NotContains(t, p.Addresses(), addr2) + }) +} + +/////////////////////////////////////////////////////////////// + +type mockManager struct { + t *testing.T + server *httptest.Server + + poller *poller.Poller + addresses map[dax.Address]struct{} +} + +func newMockManager(t *testing.T, ctx context.Context, deregisterPath string, addrs []dax.Address) *mockManager { + addresses := make(map[dax.Address]struct{}) + for _, addr := range addrs { + addresses[addr] = struct{}{} + } + + mm := &mockManager{ + t: t, + addresses: addresses, + } + + // deregister is a function used in this mock to remove the address from the + // addresses cache in the mock manager, as well as call RemoveAddresses on + // the Poller. + deregister := func(addrs ...dax.Address) { + for _, addr := range addrs { + delete(mm.addresses, addr) + } + mm.poller.RemoveAddresses(ctx, addrs...) + } + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, "/"+deregisterPath, r.URL.Path) + assert.Equal(t, "application/json", r.Header.Get("Content-Type")) + assert.Equal(t, "application/json", r.Header.Get("Accept")) + + ////// handle payload + body := r.Body + defer body.Close() + + req := mds_http.DeregisterNodesRequest{} + if err := json.NewDecoder(body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + log.Printf("deregister addresses: %s", req.Addresses) + deregister(req.Addresses...) + ////// + + w.WriteHeader(http.StatusOK) + })) + + mm.server = server + return mm +} + +func (m *mockManager) setPoller(p *poller.Poller) { + m.poller = p +} + +func (m *mockManager) URL() string { + if m.server != nil { + return m.server.URL + } + return "" +} + +func (m *mockManager) Close() { + if m.server != nil { + m.server.Close() + } +} + +type mockNode struct { + t *testing.T + server *httptest.Server +} + +func newMockNode(t *testing.T, healthPath string, dieAfter time.Duration) *mockNode { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, "/"+healthPath, r.URL.Path) + w.WriteHeader(http.StatusOK) + })) + if dieAfter > 0 { + go func() { + time.Sleep(dieAfter) + log.Printf("stopping node: %s", server.URL) + server.Close() + }() + } + return &mockNode{ + t: t, + server: server, + } +} + +func (m *mockNode) URL() string { + if m.server != nil { + return m.server.URL + } + return "" +} + +func (m *mockNode) Close() { + if m.server != nil { + m.server.Close() + } +} diff --git a/dax/mds/schemar/boltdb/schemar.go b/dax/mds/schemar/boltdb/schemar.go new file mode 100644 index 000000000..1d4ec02b2 --- /dev/null +++ b/dax/mds/schemar/boltdb/schemar.go @@ -0,0 +1,386 @@ +// Package boltdb contains the boltdb implementation of the Schemar +// interfaces. +package boltdb + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "strings" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/boltdb" + "github.com/molecula/featurebase/v3/dax/mds/schemar" + "github.com/molecula/featurebase/v3/errors" + "github.com/molecula/featurebase/v3/logger" +) + +var ( + bucketSchemar = boltdb.Bucket("schemar") +) + +// SchemarBuckets defines the buckets used by this package. It can be called +// during setup to create the buckets ahead of time. +var SchemarBuckets []boltdb.Bucket = []boltdb.Bucket{ + bucketSchemar, +} + +// Ensure type implements interface. +var _ schemar.Schemar = (*Schemar)(nil) + +type Schemar struct { + db *boltdb.DB + + logger logger.Logger +} + +// NewSchemar returns a new instance of Schemar with default values. +func NewSchemar(db *boltdb.DB, logger logger.Logger) *Schemar { + return &Schemar{ + db: db, + logger: logger, + } +} + +// CreateTable creates the table provided. If a table with the same name already +// exists then an error is returned. +func (s *Schemar) CreateTable(ctx context.Context, qtbl *dax.QualifiedTable) error { + // Ensure the table id is not blank. + if qtbl.ID == "" { + return schemar.NewErrTableIDInvalid(qtbl.ID) + } + + // Ensure the table name is not blank. + if qtbl.Name == "" { + return schemar.NewErrTableNameInvalid(qtbl.Name) + } + + // Ensure that a primary key field is present and valid. + if !qtbl.HasValidPrimaryKey() { + return schemar.NewErrInvalidPrimaryKey() + } + + //////////// end validation + + tx, err := s.db.BeginTx(ctx, true) + if err != nil { + return errors.Wrap(err, "getting transaction") + } + defer tx.Rollback() + + // Ensure a table with that ID doesn't already exist. + if t, _ := s.tableByID(tx, qtbl.TableQualifier, qtbl.ID); t != nil { + return dax.NewErrTableIDExists(qtbl.QualifiedID()) + } + + if err := s.putTable(tx, qtbl); err != nil { + return errors.Wrap(err, "putting table") + } + + // In addition to storing the table in tableKey, we want to store a reverse-lookup + // (i.e. index) on table name to the tableKey. + if err := s.putTableName(tx, qtbl); err != nil { + return errors.Wrap(err, "putting table name") + } + + return tx.Commit() +} + +// CreateField creates the field provided in the given table. If a field with +// the same name already exists then an error is returned. +func (s *Schemar) CreateField(ctx context.Context, qtid dax.QualifiedTableID, fld *dax.Field) error { + // Ensure the field name is not blank. + if fld.Name == "" { + return schemar.NewErrFieldNameInvalid(fld.Name) + } + + //////////// end validation + + tx, err := s.db.BeginTx(ctx, true) + if err != nil { + return errors.Wrap(err, "getting transaction") + } + defer tx.Rollback() + + // Get the table. + qtbl, err := s.tableByQTID(tx, qtid) + if err != nil { + return errors.Wrap(err, "getting table by id") + } + + // Ensure a field with that name doesn't already exist. + if _, ok := qtbl.Field(fld.Name); ok { + return dax.NewErrFieldExists(fld.Name) + } + + qtbl.Fields = append(qtbl.Fields, fld) + + // Write table back to database. + if err := s.putTable(tx, qtbl); err != nil { + return errors.Wrap(err, "putting table") + } + + return tx.Commit() +} + +// DropField removes the field from the table. +func (s *Schemar) DropField(ctx context.Context, qtid dax.QualifiedTableID, fldName dax.FieldName) error { + tx, err := s.db.BeginTx(ctx, true) + if err != nil { + return errors.Wrap(err, "getting transaction") + } + defer tx.Rollback() + + // Get the table. + qtbl, err := s.tableByQTID(tx, qtid) + if err != nil { + return errors.Wrap(err, "getting table by id") + } + + // Ensure a field with that name exists. + if _, ok := qtbl.Field(fldName); !ok { + return dax.NewErrFieldDoesNotExist(fldName) + } + + _ = qtbl.RemoveField(fldName) + + // Write table back to database. + if err := s.putTable(tx, qtbl); err != nil { + return errors.Wrap(err, "putting table") + } + + return tx.Commit() +} + +func (s *Schemar) putTable(tx *boltdb.Tx, qtbl *dax.QualifiedTable) error { + bkt := tx.Bucket(bucketSchemar) + if bkt == nil { + return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketSchemar) + } + + val, err := json.Marshal(qtbl) + if err != nil { + return errors.Wrap(err, "marshalling table to json") + } + + return bkt.Put(tableKey(qtbl.OrganizationID, qtbl.DatabaseID, qtbl.Table.ID), val) +} + +func (s *Schemar) putTableName(tx *boltdb.Tx, qtbl *dax.QualifiedTable) error { + bkt := tx.Bucket(bucketSchemar) + if bkt == nil { + return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketSchemar) + } + + return bkt.Put(tableNameKey(qtbl.OrganizationID, qtbl.DatabaseID, qtbl.Name), tableKey(qtbl.OrganizationID, qtbl.DatabaseID, qtbl.ID)) +} + +// Table returns the TableInfo for the given table. An error is returned if the +// table does not exist. +func (s *Schemar) Table(ctx context.Context, qtid dax.QualifiedTableID) (*dax.QualifiedTable, error) { + tx, err := s.db.BeginTx(ctx, false) + if err != nil { + return nil, errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + return s.tableByQTID(tx, qtid) +} + +// tableByQTID gets the full qualified table by the QualifiedTableID whether it has Name or ID set. +func (s *Schemar) tableByQTID(tx *boltdb.Tx, qtid dax.QualifiedTableID) (*dax.QualifiedTable, error) { + if qtid.ID == "" { + return s.tableByName(tx, qtid.TableQualifier, qtid.Name) + } + + return s.tableByID(tx, qtid.TableQualifier, qtid.ID) +} + +func (s *Schemar) tableByName(tx *boltdb.Tx, qual dax.TableQualifier, name dax.TableName) (*dax.QualifiedTable, error) { + qtid, err := s.tableIDByName(tx, qual, name) + if err != nil { + return nil, errors.Wrap(err, "getting table ID") + } + + return s.tableByID(tx, qtid.TableQualifier, qtid.ID) // TODO remove? +} + +func (s *Schemar) tableByID(tx *boltdb.Tx, qual dax.TableQualifier, id dax.TableID) (*dax.QualifiedTable, error) { + bkt := tx.Bucket(bucketSchemar) + if bkt == nil { + return nil, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketSchemar) + } + + b := bkt.Get(tableKey(qual.OrganizationID, qual.DatabaseID, id)) + if b == nil { + return nil, dax.NewErrTableIDDoesNotExist(dax.QualifiedTableID{TableQualifier: qual, ID: id}) + } + + table := &dax.QualifiedTable{} + if err := json.Unmarshal(b, table); err != nil { + return nil, errors.Wrap(err, "unmarshalling table json") + } + + return table, nil +} + +func (s *Schemar) tableIDByName(tx *boltdb.Tx, qual dax.TableQualifier, name dax.TableName) (dax.QualifiedTableID, error) { + bkt := tx.Bucket(bucketSchemar) + if bkt == nil { + return dax.QualifiedTableID{}, errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketSchemar) + } + + b := bkt.Get(tableNameKey(qual.OrganizationID, qual.DatabaseID, name)) + if b == nil { + return dax.QualifiedTableID{}, dax.NewErrTableNameDoesNotExist(name) + } + + return keyQualifiedTableID(b) +} + +// Tables returns a list of Table for all existing tables. If one or more table +// names is provided, then only those will be included in the output. +func (s *Schemar) Tables(ctx context.Context, qual dax.TableQualifier, ids ...dax.TableID) ([]*dax.QualifiedTable, error) { + tx, err := s.db.BeginTx(ctx, false) + if err != nil { + return nil, errors.Wrap(err, "beginning tx") + } + defer tx.Rollback() + + return s.getTables(ctx, tx, qual, ids...) +} + +func (s *Schemar) getTables(ctx context.Context, tx *boltdb.Tx, qual dax.TableQualifier, ids ...dax.TableID) (dax.QualifiedTables, error) { + c := tx.Bucket(bucketSchemar).Cursor() + + // Deserialize rows into Table objects. + tables := make(dax.QualifiedTables, 0) + + var filterByID bool + if len(ids) > 0 { + filterByID = true + } + + prefix := []byte(fmt.Sprintf(prefixFmtTables, qual.OrganizationID, qual.DatabaseID)) + for k, v := c.Seek(prefix); k != nil && bytes.HasPrefix(k, prefix); k, v = c.Next() { + if v == nil { + s.logger.Printf("nil value for key: %s", k) + continue + } + + tblID, err := keyTableID(k) + if err != nil { + return nil, errors.Wrap(err, "getting table from key") + } + + // Only include tables provided in the ids filter. + if filterByID && !containsTableID(ids, tblID) { + continue + } + + table := &dax.QualifiedTable{} + if err := json.Unmarshal(v, table); err != nil { + return nil, errors.Wrap(err, "unmarshalling table json") + } + + tables = append(tables, table) + } + + return tables, nil +} + +func containsTableID(s []dax.TableID, e dax.TableID) bool { + for _, a := range s { + if a == e { + return true + } + } + return false +} + +// DropTable drops the given table. If the named/IDed table does not exist +// then an error is returned. +func (s *Schemar) DropTable(ctx context.Context, qtid dax.QualifiedTableID) error { + tx, err := s.db.BeginTx(ctx, true) + if err != nil { + return errors.Wrap(err, "getting transaction") + } + defer tx.Rollback() + + // Ensure the table exists. + qtbl, err := s.tableByQTID(tx, qtid) + if err != nil { + return errors.Wrap(err, "getting table by id") + } + + bkt := tx.Bucket(bucketSchemar) + if bkt == nil { + return errors.Errorf(boltdb.ErrFmtBucketNotFound, bucketSchemar) + } + + // Delete the table by ID. + if err := bkt.Delete(tableKey(qtbl.OrganizationID, qtbl.DatabaseID, qtbl.ID)); err != nil { + return errors.Wrap(err, "deleting table by id") + } + + // Delete the reverse-lookup table by Name. + if err := bkt.Delete(tableNameKey(qtbl.OrganizationID, qtbl.DatabaseID, qtbl.Name)); err != nil { + return errors.Wrap(err, "deleting table by name") + } + + return tx.Commit() +} + +const ( + prefixFmtTables = "tables/%s/%s/" + prefixFmtTableNames = "tablenames/%s/%s/" +) + +// tableKey returns a key based on a qualified table ID. +func tableKey(orgID dax.OrganizationID, dbID dax.DatabaseID, tblID dax.TableID) []byte { + key := fmt.Sprintf(prefixFmtTables+"%s", orgID, dbID, tblID) + return []byte(key) +} + +// tableNameKey returns a key based on a qualified table name. +func tableNameKey(orgID dax.OrganizationID, dbID dax.DatabaseID, name dax.TableName) []byte { + key := fmt.Sprintf(prefixFmtTableNames+"%s", orgID, dbID, name) + return []byte(key) +} + +// keyTableID gets the TableID out of the key. +func keyTableID(key []byte) (dax.TableID, error) { + parts := strings.Split(string(key), "/") + if len(parts) != 4 { + return "", errors.New(errors.ErrUncoded, "table key format expected: `tables/orgID/dbID/tblID`") + } + + return dax.TableID(parts[3]), nil +} + +// keyQualifedTableID gets the QualifiedTableID out of the key. +func keyQualifiedTableID(key []byte) (dax.QualifiedTableID, error) { + parts := strings.Split(string(key), "/") + if len(parts) != 4 { + return dax.QualifiedTableID{}, errors.New(errors.ErrUncoded, "table key format expected: `tables/orgID/dbID/tblID`") + } + + return dax.NewQualifiedTableID( + dax.NewTableQualifier( + dax.OrganizationID(parts[1]), + dax.DatabaseID(parts[2]), + ), + dax.TableID(parts[3]), + ), nil +} + +func (s *Schemar) TableID(ctx context.Context, qual dax.TableQualifier, name dax.TableName) (dax.QualifiedTableID, error) { + tx, err := s.db.BeginTx(ctx, false) + if err != nil { + return dax.QualifiedTableID{}, err + } + defer tx.Rollback() + + return s.tableIDByName(tx, qual, name) +} diff --git a/dax/mds/schemar/boltdb/schemar_test.go b/dax/mds/schemar/boltdb/schemar_test.go new file mode 100644 index 000000000..4460e7d66 --- /dev/null +++ b/dax/mds/schemar/boltdb/schemar_test.go @@ -0,0 +1,146 @@ +package boltdb_test + +import ( + "context" + "testing" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/mds/schemar/boltdb" + daxtest "github.com/molecula/featurebase/v3/dax/test" + testbolt "github.com/molecula/featurebase/v3/dax/test/boltdb" + "github.com/molecula/featurebase/v3/errors" + "github.com/molecula/featurebase/v3/logger" + "github.com/stretchr/testify/assert" +) + +func TestSchemar(t *testing.T) { + orgID := dax.OrganizationID("acme") + dbID := dax.DatabaseID("db1") + invalidTableID := dax.TableID("invalidID") + tableName := dax.TableName("foo") + tableName0 := dax.TableName("foo") + tableName1 := dax.TableName("bar") + tableID0 := "2" + tableID1 := "1" + partitionN := 12 + + ctx := context.Background() + qual := dax.NewTableQualifier(orgID, dbID) + + db := testbolt.MustOpenDB(t) + defer testbolt.MustCloseDB(t, db) + + t.Cleanup(func() { + testbolt.CleanupDB(t, db.Path()) + }) + + // Initialize the buckets. + assert.NoError(t, db.InitializeBuckets(boltdb.SchemarBuckets...)) + + t.Run("NewSchemar", func(t *testing.T) { + s := boltdb.NewSchemar(db, logger.NopLogger) + + // Add new table. + tbl := dax.NewTable(tableName) + tbl.CreateID() + tbl.Fields = []*dax.Field{ + { + Name: dax.PrimaryKeyFieldName, + Type: dax.FieldTypeString, + }, + { + Name: "intField", + Type: dax.FieldTypeInt, + }, + } + qtbl := dax.NewQualifiedTable(qual, tbl) + assert.NoError(t, s.CreateTable(ctx, qtbl)) + + // Try adding the table again. + err := s.CreateTable(ctx, qtbl) + if assert.Error(t, err) { + assert.True(t, errors.Is(err, dax.ErrTableIDExists)) + } + + qtid := qtbl.QualifiedID() + + // Get the table. + { + tbl, err := s.Table(ctx, qtid) + assert.NoError(t, err) + assert.Equal(t, tableName, tbl.Name) + } + + // Drop the table. + assert.NoError(t, s.DropTable(ctx, qtid)) + + // Make sure the reverse-lookup (table by name) was dropped as well. + { + _, err := s.TableID(ctx, qual, tableName) + if assert.Error(t, err) { + assert.True(t, errors.Is(err, dax.ErrTableNameDoesNotExist)) + } + } + + // Try adding the table (i.e. the same table name) again. + assert.NoError(t, s.CreateTable(ctx, qtbl)) + + // Drop the table again. + assert.NoError(t, s.DropTable(ctx, qtid)) + + // Drop invalid table. + { + iqtid := dax.NewQualifiedTableID(qual, invalidTableID) + err := s.DropTable(ctx, iqtid) + if assert.Error(t, err) { + assert.True(t, errors.Is(err, dax.ErrTableIDDoesNotExist)) + } + } + }) + + t.Run("GetTables", func(t *testing.T) { + s := boltdb.NewSchemar(db, logger.NopLogger) + + exp := []*dax.QualifiedTable{} + tables, err := s.Tables(ctx, qual) + assert.NoError(t, err) + assert.Equal(t, exp, tables) + + qtbl0 := daxtest.TestQualifiedTableWithID(t, qual, tableID0, tableName0, partitionN, false) + qtbl1 := daxtest.TestQualifiedTableWithID(t, qual, tableID1, tableName1, partitionN, false) + + // Add a couple of tables. + assert.NoError(t, s.CreateTable(ctx, qtbl0)) + assert.NoError(t, s.CreateTable(ctx, qtbl1)) + + exp = []*dax.QualifiedTable{ + qtbl1, + qtbl0, + } + + // All tables. + tables, err = s.Tables(ctx, qual) + assert.NoError(t, err) + assert.Equal(t, exp, tables) + + // With a valid filter. + tables, err = s.Tables(ctx, qual, qtbl0.ID) + assert.NoError(t, err) + assert.Equal(t, exp[1:], tables) + + // With an invalid filter. + tables, err = s.Tables(ctx, qual, invalidTableID) + assert.NoError(t, err) + assert.Equal(t, exp[0:0], tables) + + // With both valid and invalid filters. + tables, err = s.Tables(ctx, qual, qtbl0.ID, invalidTableID) + assert.NoError(t, err) + assert.Equal(t, exp[1:], tables) + + // With all valid filters. + tables, err = s.Tables(ctx, qual, qtbl0.ID, qtbl1.ID) + assert.NoError(t, err) + assert.Equal(t, exp, tables) + }) +} diff --git a/dax/mds/schemar/errors.go b/dax/mds/schemar/errors.go new file mode 100644 index 000000000..12ed61aad --- /dev/null +++ b/dax/mds/schemar/errors.go @@ -0,0 +1,44 @@ +package schemar + +import ( + "fmt" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/errors" +) + +const ( + ErrCodeTableIDInvalid errors.Code = "TableIDInvalid" + ErrCodeTableNameInvalid errors.Code = "TableNameInvalid" + ErrCodeInvalidPrimaryKey errors.Code = "InvalidPrimaryKey" + + ErrCodeFieldNameInvalid errors.Code = "FieldNameInvalid" +) + +func NewErrTableIDInvalid(tableID dax.TableID) error { + return errors.New( + ErrCodeTableIDInvalid, + fmt.Sprintf("table ID '%s' is invalid", tableID), + ) +} + +func NewErrTableNameInvalid(tableName dax.TableName) error { + return errors.New( + ErrCodeTableNameInvalid, + fmt.Sprintf("table name '%s' is invalid", tableName), + ) +} + +func NewErrInvalidPrimaryKey() error { + return errors.New( + ErrCodeInvalidPrimaryKey, + "invalid primary key", + ) +} + +func NewErrFieldNameInvalid(fieldName dax.FieldName) error { + return errors.New( + ErrCodeFieldNameInvalid, + fmt.Sprintf("field name '%s' is invalid", fieldName), + ) +} diff --git a/dax/mds/schemar/http/handler.go b/dax/mds/schemar/http/handler.go new file mode 100644 index 000000000..18e6969ef --- /dev/null +++ b/dax/mds/schemar/http/handler.go @@ -0,0 +1,199 @@ +package http + +import ( + "encoding/json" + "net/http" + + "github.com/gorilla/mux" + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/mds/schemar" +) + +func Handler(s schemar.Schemar) http.Handler { + svr := &server{ + schemar: s, + } + + router := mux.NewRouter() + router.HandleFunc("/health", svr.getHealth).Methods("GET").Name("GetHealth") + router.HandleFunc("/create-table", svr.postCreateTable).Methods("POST").Name("PostCreateTable") + router.HandleFunc("/drop-table", svr.postDropTable).Methods("POST").Name("PostDropTable") + router.HandleFunc("/table", svr.postTable).Methods("POST").Name("PostTable") + router.HandleFunc("/tables", svr.postTables).Methods("POST").Name("PostTables") + return router +} + +type server struct { + schemar schemar.Schemar +} + +// GET /health +func (s *server) getHealth(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) +} + +// POST /create-table +func (s *server) postCreateTable(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + ctx := r.Context() + + req := &dax.QualifiedTable{} + if err := json.NewDecoder(body).Decode(req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + err := s.schemar.CreateTable(ctx, req) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + resp := struct{}{} + if err := json.NewEncoder(w).Encode(resp); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } +} + +// POST /drop-table +func (s *server) postDropTable(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + ctx := r.Context() + + req := DropTableRequest{} + if err := json.NewDecoder(body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + qtid := req.TableKey.QualifiedTableID() + + err := s.schemar.DropTable(ctx, qtid) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + resp := struct{}{} + if err := json.NewEncoder(w).Encode(resp); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } +} + +// // POST /create-field +// func (s *server) postCreateField(w http.ResponseWriter, r *http.Request) { +// body := r.Body +// defer body.Close() + +// req := mds.CreateFieldRequest{} +// if err := json.NewDecoder(body).Decode(&req); err != nil { +// http.Error(w, err.Error(), http.StatusBadRequest) +// return +// } + +// resp, err := s.mds.CreateField(req) +// if err != nil { +// http.Error(w, err.Error(), http.StatusBadRequest) +// return +// } + +// if err := json.NewEncoder(w).Encode(resp); err != nil { +// http.Error(w, err.Error(), http.StatusBadRequest) +// return +// } +// } + +// // POST /drop-field +// func (s *server) postDropField(w http.ResponseWriter, r *http.Request) { +// body := r.Body +// defer body.Close() + +// req := mds.DropFieldRequest{} +// if err := json.NewDecoder(body).Decode(&req); err != nil { +// http.Error(w, err.Error(), http.StatusBadRequest) +// return +// } + +// resp, err := s.mds.DropField(req) +// if err != nil { +// http.Error(w, err.Error(), http.StatusBadRequest) +// return +// } + +// if err := json.NewEncoder(w).Encode(resp); err != nil { +// http.Error(w, err.Error(), http.StatusBadRequest) +// return +// } +// } + +// POST /table +func (s *server) postTable(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + ctx := r.Context() + + req := TableRequest{} + if err := json.NewDecoder(body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + qtid := req.TableKey.QualifiedTableID() + + resp, err := s.schemar.Table(ctx, qtid) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + if err := json.NewEncoder(w).Encode(resp); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } +} + +// POST /tables +func (s *server) postTables(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + ctx := r.Context() + + req := TablesRequest{} + if err := json.NewDecoder(body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + qual := dax.NewTableQualifier(req.OrganizationID, req.DatabaseID) + resp, err := s.schemar.Tables(ctx, qual, req.TableIDs...) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + if err := json.NewEncoder(w).Encode(resp); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } +} + +type DropTableRequest struct { + TableKey dax.TableKey `json:"table-key"` +} + +type TableRequest struct { + TableKey dax.TableKey `json:"table-key"` +} + +type TablesRequest struct { + OrganizationID dax.OrganizationID `json:"org-id"` + DatabaseID dax.DatabaseID `json:"db-id"` + TableIDs dax.TableIDs `json:"table-ids"` +} diff --git a/dax/mds/schemar/schemar.go b/dax/mds/schemar/schemar.go new file mode 100644 index 000000000..b1639f137 --- /dev/null +++ b/dax/mds/schemar/schemar.go @@ -0,0 +1,54 @@ +// Package schemar provides the core Schemar interface. +package schemar + +import ( + "context" + + "github.com/molecula/featurebase/v3/dax" +) + +type Schemar interface { + CreateTable(context.Context, *dax.QualifiedTable) error + DropTable(context.Context, dax.QualifiedTableID) error + CreateField(context.Context, dax.QualifiedTableID, *dax.Field) error + DropField(context.Context, dax.QualifiedTableID, dax.FieldName) error + Table(context.Context, dax.QualifiedTableID) (*dax.QualifiedTable, error) + Tables(context.Context, dax.TableQualifier, ...dax.TableID) ([]*dax.QualifiedTable, error) + + // TableID is a reverse-lookup method to get the TableID for a given + // qualified TableName. + TableID(context.Context, dax.TableQualifier, dax.TableName) (dax.QualifiedTableID, error) +} + +////////////////////////////////////////////// + +// Ensure type implements interface. +var _ Schemar = &NopSchemar{} + +// NopSchemar is a no-op implementation of the Schemar interface. +type NopSchemar struct{} + +func NewNopSchemar() *NopSchemar { + return &NopSchemar{} +} + +func (s *NopSchemar) CreateTable(ctx context.Context, qtbl *dax.QualifiedTable) error { return nil } +func (s *NopSchemar) DropTable(ctx context.Context, qtid dax.QualifiedTableID) error { + return nil +} +func (s *NopSchemar) CreateField(ctx context.Context, qtid dax.QualifiedTableID, fld *dax.Field) error { + return nil +} +func (s *NopSchemar) DropField(ctx context.Context, qtid dax.QualifiedTableID, fld dax.FieldName) error { + return nil +} +func (s *NopSchemar) Table(ctx context.Context, qtid dax.QualifiedTableID) (*dax.QualifiedTable, error) { + return nil, nil +} +func (s *NopSchemar) Tables(ctx context.Context, qual dax.TableQualifier, ids ...dax.TableID) ([]*dax.QualifiedTable, error) { + return []*dax.QualifiedTable{}, nil +} + +func (s *NopSchemar) TableID(context.Context, dax.TableQualifier, dax.TableName) (dax.QualifiedTableID, error) { + return dax.QualifiedTableID{}, nil +} diff --git a/dax/mds/schemar/schemar_test.go b/dax/mds/schemar/schemar_test.go new file mode 100644 index 000000000..62a4f2edf --- /dev/null +++ b/dax/mds/schemar/schemar_test.go @@ -0,0 +1,124 @@ +package schemar_test + +import ( + "context" + "testing" + + "github.com/molecula/featurebase/v3/dax" + daxtest "github.com/molecula/featurebase/v3/dax/test" + "github.com/molecula/featurebase/v3/errors" + "github.com/stretchr/testify/assert" +) + +func TestSchemar(t *testing.T) { + orgID := dax.OrganizationID("acme") + dbID := dax.DatabaseID("db1") + invalidTableID := dax.TableID("invalidID") + tableName := dax.TableName("foo") + tableName0 := dax.TableName("foo") + tableName1 := dax.TableName("bar") + tableID0 := "2" + tableID1 := "1" + partitionN := 12 + + ctx := context.Background() + qual := dax.NewTableQualifier(orgID, dbID) + + t.Run("NewSchemar", func(t *testing.T) { + s, cleanup := daxtest.NewSchemar(t) + defer cleanup() + + // Add new table. + tbl := dax.NewTable(tableName) + tbl.Fields = []*dax.Field{ + { + Name: dax.PrimaryKeyFieldName, + Type: dax.FieldTypeString, + }, + { + Name: "intField", + Type: dax.FieldTypeInt, + }, + } + qtbl := dax.NewQualifiedTable(qual, tbl) + qtbl.CreateID() + assert.NoError(t, s.CreateTable(ctx, qtbl)) + + // Try adding the table again. + err := s.CreateTable(ctx, qtbl) + if assert.Error(t, err) { + assert.True(t, errors.Is(err, dax.ErrTableIDExists)) + } + + qtid := qtbl.QualifiedID() + + // Get the table. + { + tbl, err := s.Table(ctx, qtid) + assert.NoError(t, err) + assert.Equal(t, tableName, tbl.Name) + } + + // Drop the table. + { + err := s.DropTable(ctx, qtid) + assert.NoError(t, err) + } + + // Drop invalid table. + { + iqtid := dax.NewQualifiedTableID(qual, invalidTableID) + err := s.DropTable(ctx, iqtid) + if assert.Error(t, err) { + assert.True(t, errors.Is(err, dax.ErrTableIDDoesNotExist)) + } + } + }) + + t.Run("GetTables", func(t *testing.T) { + s, cleanup := daxtest.NewSchemar(t) + defer cleanup() + + exp := []*dax.QualifiedTable{} + tables, err := s.Tables(ctx, qual) + assert.NoError(t, err) + assert.Equal(t, exp, tables) + + qtbl0 := daxtest.TestQualifiedTableWithID(t, qual, tableID0, tableName0, partitionN, false) + qtbl1 := daxtest.TestQualifiedTableWithID(t, qual, tableID1, tableName1, partitionN, false) + + // Add a couple of tables. + assert.NoError(t, s.CreateTable(ctx, qtbl0)) + assert.NoError(t, s.CreateTable(ctx, qtbl1)) + + exp = []*dax.QualifiedTable{ + qtbl1, + qtbl0, + } + + // All tables. + tables, err = s.Tables(ctx, qual) + assert.NoError(t, err) + assert.Equal(t, exp, tables) + + // With a valid filter. + tables, err = s.Tables(ctx, qual, qtbl0.ID) + assert.NoError(t, err) + assert.Equal(t, exp[1:], tables) + + // With an invalid filter. + tables, err = s.Tables(ctx, qual, invalidTableID) + assert.NoError(t, err) + assert.Equal(t, exp[0:0], tables) + + // With both valid and invalid filters. + tables, err = s.Tables(ctx, qual, qtbl0.ID, invalidTableID) + assert.NoError(t, err) + assert.Equal(t, exp[1:], tables) + + // With all valid filters. + tables, err = s.Tables(ctx, qual, qtbl0.ID, qtbl1.ID) + assert.NoError(t, err) + assert.Equal(t, exp, tables) + }) +} diff --git a/dax/node.go b/dax/node.go new file mode 100644 index 000000000..8b6f9d331 --- /dev/null +++ b/dax/node.go @@ -0,0 +1,45 @@ +package dax + +import ( + "context" + "fmt" + + "github.com/molecula/featurebase/v3/errors" +) + +// Node is used in API requests, like RegisterNode (before being assigned +// roles). +type Node struct { + Address Address `json:"address"` + + RoleTypes []RoleType `json:"role-types"` +} + +// AssignedNode is used in API responses. +type AssignedNode struct { + Address Address `json:"address"` + Role Role `json:"role"` +} + +// NodeService represents a service for managing Nodes. +type NodeService interface { + CreateNode(context.Context, Address, *Node) error + ReadNode(context.Context, Address) (*Node, error) + DeleteNode(context.Context, Address) error + Nodes(context.Context) ([]*Node, error) +} + +//////////////////////////////////////////////////// +// Errors +//////////////////////////////////////////////////// + +const ( + ErrNodeDoesNotExist errors.Code = "NodeDoesNotExist" +) + +func NewErrNodeDoesNotExist(addr Address) error { + return errors.New( + ErrNodeDoesNotExist, + fmt.Sprintf("node '%s' does not exist", addr), + ) +} diff --git a/dax/partition.go b/dax/partition.go new file mode 100644 index 000000000..b1c4e297c --- /dev/null +++ b/dax/partition.go @@ -0,0 +1,65 @@ +package dax + +import "fmt" + +// PartitionNum is the numerical (int) partition value. +type PartitionNum int + +// PartitionNums is a slice of PartitionNum. +type PartitionNums []PartitionNum + +// String returns the PartitionNum as a string. +func (p PartitionNum) String() string { + return fmt.Sprintf("%d", p) +} + +// Partition is a versioned partition. +type Partition struct { + Num PartitionNum `json:"num"` + Version int `json:"version"` +} + +// NewPartition returns a Partition with the provided num and version. +func NewPartition(num PartitionNum, version int) Partition { + return Partition{ + Num: num, + Version: version, + } +} + +// String returns the Partition (i.e. its Num and Version) as a string. +func (p Partition) String() string { + return fmt.Sprintf("%d.%d", p.Num, p.Version) +} + +// Partitions is a sortable slice of Partition. +type Partitions []Partition + +func (p Partitions) Len() int { return len(p) } +func (p Partitions) Less(i, j int) bool { return p[i].Num < p[j].Num } +func (p Partitions) Swap(i, j int) { p[i], p[j] = p[j], p[i] } + +// NewPartitions returns the provided list of partition nums as a list of +// Partition with an invalid version (-1). This is to use for cases where the +// request should not be aware of a partition versioning. +func NewPartitions(partitionNums ...PartitionNum) Partitions { + pvs := make(Partitions, len(partitionNums)) + + for i := range partitionNums { + pvs[i] = Partition{ + Num: partitionNums[i], + Version: -1, + } + } + + return pvs +} + +// Nums returns a slice of all the partition numbers in Partitions. +func (p Partitions) Nums() []PartitionNum { + pp := make([]PartitionNum, len(p)) + for i := range p { + pp[i] = p[i].Num + } + return pp +} diff --git a/dax/queryer/alpha/router.go b/dax/queryer/alpha/router.go new file mode 100644 index 000000000..57362ef06 --- /dev/null +++ b/dax/queryer/alpha/router.go @@ -0,0 +1,36 @@ +package alpha + +import ( + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/queryer" + "github.com/molecula/featurebase/v3/errors" + featurebaseserver "github.com/molecula/featurebase/v3/server" +) + +// Ensure type implements interface. +var _ queryer.Router = (*Router)(nil) + +type Router struct { + computers map[dax.Address]*featurebaseserver.Command +} + +func NewRouter() *Router { + return &Router{ + computers: make(map[dax.Address]*featurebaseserver.Command), + } +} + +func (r *Router) AddCmd(addr dax.Address, cmd *featurebaseserver.Command) error { + if cmd == nil { + return errors.New(errors.ErrUncoded, "cannot add nil cmd to director") + } + r.computers[addr] = cmd + return nil +} + +func (r *Router) Importer(addr dax.Address) queryer.Importer { + if cmd, found := r.computers[addr]; found { + return queryer.NewFeatureBaseImporter(cmd.API) + } + return nil +} diff --git a/dax/queryer/api/openapi.yaml b/dax/queryer/api/openapi.yaml new file mode 100644 index 000000000..7aa748ac2 --- /dev/null +++ b/dax/queryer/api/openapi.yaml @@ -0,0 +1,99 @@ +openapi: 3.0.3 + +info: + title: Queryer + description: The query layer of the DAX architecture. + version: 0.0.0 + +paths: + /queryer/health: + get: + summary: Health check endpoint. + description: Provides an endpoint to check the overall health of the Queryer service. + operationId: GetHealth + responses: + 200: + description: Service is healthy. + + /queryer/query: + post: + summary: Execute either a PQL or SQL command. + description: Executes the given PQL or SQL command based on input, and returns the results in a standard format. + operationId: PostQuery + requestBody: + content: + application/json: + examples: + pql: + summary: Query via PQL + value: + table: tbl + pql: Row(fld=1) + sql: + summary: Query via SQL + value: + sql: SELECT * from tbl + schema: + type: object + properties: + table: + type: string + pql: + type: string + sql: + type: string + responses: + 200: + $ref: '#/components/responses/QueryResponse' + + /queryer/sql: + post: + summary: Execute a SQL command. + description: Executes the given SQL command, and returns the results in a standard format. + operationId: PostSQL + requestBody: + content: + text/plain: + example: SELECT * FROM tbl + schema: + type: string + responses: + 200: + $ref: '#/components/responses/QueryResponse' + +components: + responses: + QueryResponse: + description: Standard tabular response with optional error and warnings. + content: + application/json: + schema: + $ref: '#/components/schemas/QueryResult' + schemas: + QueryResult: + type: object + properties: + schema: + type: array # fields + items: + type: object # field + properties: + name: + type: string # column name + type: + type: string # column type + data: + type: array + items: + type: array + items: + type: string # this could really be any type; interface{} + error: + type: string + warnings: + type: array + items: + type: string + exec_time: + type: integer + format: int64 \ No newline at end of file diff --git a/dax/queryer/compute_api.go b/dax/queryer/compute_api.go new file mode 100644 index 000000000..32ac3ac85 --- /dev/null +++ b/dax/queryer/compute_api.go @@ -0,0 +1,329 @@ +package queryer + +import ( + "context" + "time" + + featurebase "github.com/molecula/featurebase/v3" + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/mds/controller/partitioner" + "github.com/molecula/featurebase/v3/errors" +) + +// Ensure type implements interface. +var _ featurebase.ComputeAPI = &qualifiedComputeAPI{} + +type qualifiedComputeAPI struct { + mds MDS + router Router + qual dax.TableQualifier +} + +func NewQualifiedComputeAPI(qual dax.TableQualifier, mds MDS, router Router) *qualifiedComputeAPI { + c := &qualifiedComputeAPI{ + mds: mds, + router: NewNopRouter(), + qual: qual, + } + + if router != nil { + c.router = router + } + + return c +} + +// importer is use to get the Importer based on the provided address. If the +// computeAPI has been configured with entries in an ImporterRouter (which is a +// map of dax.Address to in-process compute API), then it will use that. +// Otherwise, it sets up an http client based on the provided address. +func (c *qualifiedComputeAPI) importer(addr dax.Address) (Importer, error) { + if imp := c.router.Importer(addr); imp != nil { + return imp, nil + } + + return NewComputeImporter(addr), nil +} + +func (c *qualifiedComputeAPI) Import(ctx context.Context, qcx *featurebase.Qcx, req *featurebase.ImportRequest, opts ...featurebase.ImportOption) error { + // If the request is empty, return early. + if len(req.ColumnKeys) == 0 && len(req.ColumnIDs) == 0 { + return nil + } + + // Determine if the columns use string keys or not. + var hasColKeys bool + if len(req.ColumnKeys) > 0 { + hasColKeys = true + if len(req.ColumnIDs) > 0 { + return errors.Errorf("import request has both column ids and keys") + } + } + // Determine if the rows use string keys or not. + var hasRowKeys bool + if len(req.RowKeys) > 0 { + hasRowKeys = true + if len(req.RowIDs) > 0 { + return errors.Errorf("import request has both row ids and keys") + } + } + + partitioner := partitioner.NewPartitioner() + + reqPerShard := make(map[uint64]*featurebase.ImportRequest) + + tkey, err := c.indexToQualifiedTableKey(ctx, req.Index) + if err != nil { + return errors.Wrap(err, "converting index to qualified table key") + } + stkey := string(tkey) + + qtid := tkey.QualifiedTableID() + + qtbl, err := c.mds.Table(ctx, qtid) + if err != nil { + return errors.Wrapf(err, "getting table for import: %s", req.Index) + } + + // Translate column keys. + if hasColKeys { + // Get the partitions (and therefore, nodes) responsible for the keys. + pMap := partitioner.PartitionsForKeys(qtbl.Key(), qtbl.PartitionN, req.ColumnKeys...) + + colIDs := make([]uint64, 0, len(req.ColumnKeys)) + for pNum := range pMap { + addr, err := c.mds.IngestPartition(ctx, qtid, pNum) + if err != nil { + return errors.Wrapf(err, "getting ingest partition: %d", pNum) + } + + importer, err := c.importer(addr) + if err != nil { + return errors.Wrapf(err, "getting importer for address: %s", addr) + } + + colKeyMap, err := importer.CreateIndexKeys(ctx, stkey, req.ColumnKeys...) + if err != nil { + return errors.Wrap(err, "creating index keys") + } + for i := range req.ColumnKeys { + colIDs = append(colIDs, colKeyMap[req.ColumnKeys[i]]) + } + } + req.ColumnIDs = colIDs + } + + // Translate row keys. + if hasRowKeys { + addr, err := c.mds.IngestPartition(ctx, qtid, 0) + if err != nil { + return errors.Wrapf(err, "getting ingest partition: %d", 0) + } + + importer, err := c.importer(addr) + if err != nil { + return errors.Wrapf(err, "getting importer for address: %s", addr) + } + + rowKeyMap, err := importer.CreateFieldKeys(ctx, stkey, req.Field, req.RowKeys...) + if err != nil { + return errors.Wrap(err, "creating field keys") + } + rowIDs := make([]uint64, len(req.RowKeys)) + for i := range req.RowKeys { + rowIDs[i] = rowKeyMap[req.RowKeys[i]] + } + req.RowIDs = rowIDs + } + + // Loop over the column ids and split them up by shard. + for ii := range req.ColumnIDs { + // Determine shard. + shard := req.ColumnIDs[ii] / featurebase.ShardWidth + + // Get or create the ImportRequest for this shard. + shardedReq, found := reqPerShard[shard] + if !found { + shardedReq = &featurebase.ImportRequest{ + Index: stkey, + IndexCreatedAt: req.IndexCreatedAt, + Field: req.Field, + FieldCreatedAt: req.FieldCreatedAt, + Shard: shard, + RowIDs: []uint64{}, + ColumnIDs: []uint64{}, + RowKeys: []string{}, + ColumnKeys: []string{}, + Timestamps: []int64{}, + Clear: req.Clear, + } + reqPerShard[shard] = shardedReq + } + + shardedReq.ColumnIDs = append(shardedReq.ColumnIDs, req.ColumnIDs[ii]) + if len(req.RowIDs) > 0 { + shardedReq.RowIDs = append(shardedReq.RowIDs, req.RowIDs[ii]) + } + if len(req.Timestamps) > 0 { + shardedReq.Timestamps = append(shardedReq.Timestamps, req.Timestamps[ii]) + } + } + + // Send each of the sharded ImportRequests to the appropriate compute node. + for shard, req := range reqPerShard { + addr, err := c.mds.IngestShard(ctx, qtid, dax.ShardNum(shard)) + if err != nil { + return errors.Wrapf(err, "getting ingest shard: %d", shard) + } + + importer, err := c.importer(addr) + if err != nil { + return errors.Wrapf(err, "getting importer for address: %s", addr) + } + + importer.Import(ctx, req, + featurebase.OptImportOptionsClear(req.Clear), + featurebase.OptImportOptionsIgnoreKeyCheck(true), + ) + } + + return nil +} + +func (c *qualifiedComputeAPI) ImportValue(ctx context.Context, qcx *featurebase.Qcx, req *featurebase.ImportValueRequest, opts ...featurebase.ImportOption) error { + // If the request is empty, return early. + if len(req.ColumnKeys) == 0 && len(req.ColumnIDs) == 0 { + return nil + } + + // Determine if the columns use string keys or not. + var hasColKeys bool + if len(req.ColumnKeys) > 0 { + hasColKeys = true + if len(req.ColumnIDs) > 0 { + return errors.Errorf("import value request has both column ids and keys") + } + } + + partitioner := partitioner.NewPartitioner() + + reqPerShard := make(map[uint64]*featurebase.ImportValueRequest) + + tkey, err := c.indexToQualifiedTableKey(ctx, req.Index) + if err != nil { + return errors.Wrap(err, "converting index to qualified table key") + } + stkey := string(tkey) + + qtid := tkey.QualifiedTableID() + + qtbl, err := c.mds.Table(ctx, qtid) + if err != nil { + return errors.Wrapf(err, "getting table for importvalue: %s", req.Index) + } + + // Translate column keys. + if hasColKeys { + // Get the partitions (and therefore, nodes) responsible for the keys. + pMap := partitioner.PartitionsForKeys(qtbl.Key(), qtbl.PartitionN, req.ColumnKeys...) + + colIDs := make([]uint64, 0, len(req.ColumnKeys)) + for pNum := range pMap { + addr, err := c.mds.IngestPartition(ctx, qtid, pNum) + if err != nil { + return errors.Wrapf(err, "getting ingest partition: %d", pNum) + } + + importer, err := c.importer(addr) + if err != nil { + return errors.Wrapf(err, "getting importer for address: %s", addr) + } + + colKeyMap, err := importer.CreateIndexKeys(ctx, stkey, req.ColumnKeys...) + if err != nil { + return errors.Wrap(err, "creating index keys") + } + for i := range req.ColumnKeys { + colIDs = append(colIDs, colKeyMap[req.ColumnKeys[i]]) + } + } + req.ColumnIDs = colIDs + } + + // Loop over the column ids and split them up by shard. + for ii := range req.ColumnIDs { + // Determine shard. + shard := req.ColumnIDs[ii] / featurebase.ShardWidth + + // Get or create the ImportRequest for this shard. + shardedReq, found := reqPerShard[shard] + if !found { + shardedReq = &featurebase.ImportValueRequest{ + Index: stkey, + IndexCreatedAt: req.IndexCreatedAt, + Field: req.Field, + FieldCreatedAt: req.FieldCreatedAt, + Shard: shard, + ColumnIDs: []uint64{}, + ColumnKeys: []string{}, + Values: []int64{}, + FloatValues: []float64{}, + TimestampValues: []time.Time{}, + StringValues: []string{}, + Clear: req.Clear, + } + reqPerShard[shard] = shardedReq + } + + shardedReq.ColumnIDs = append(shardedReq.ColumnIDs, req.ColumnIDs[ii]) + if len(req.Values) > 0 { + shardedReq.Values = append(shardedReq.Values, req.Values[ii]) + } + // TODO: The following would populate the other value types, but the + // EncodeImportValues doesn't seem to use this data. So we need to track + // down how this is being used. + // + // if len(req.FloatValues) > 0 { + // shardedReq.FloatValues = append(shardedReq.FloatValues, req.FloatValues[ii]) + // } + // if len(req.TimestampValues) > 0 { + // shardedReq.TimestampValues = append(shardedReq.TimestampValues, req.TimestampValues[ii]) + // } + // if len(req.StringValues) > 0 { + // shardedReq.StringValues = append(shardedReq.StringValues, req.StringValues[ii]) + // } + } + + // Send each of the sharded ImportValueRequests to the appropriate compute + // node. + for shard, req := range reqPerShard { + addr, err := c.mds.IngestShard(ctx, qtid, dax.ShardNum(shard)) + if err != nil { + return errors.Wrapf(err, "getting ingest shard: %d", shard) + } + + importer, err := c.importer(addr) + if err != nil { + return errors.Wrapf(err, "getting importer for address: %s", addr) + } + + importer.ImportValue(ctx, req, + featurebase.OptImportOptionsClear(req.Clear), + featurebase.OptImportOptionsIgnoreKeyCheck(true), + ) + } + + return nil +} + +func (c *qualifiedComputeAPI) Txf() *featurebase.TxFactory { + return &featurebase.TxFactory{} +} + +func (c *qualifiedComputeAPI) indexToQualifiedTableKey(ctx context.Context, index string) (dax.TableKey, error) { + qtid, err := c.mds.TableID(ctx, c.qual, dax.TableName(index)) + if err != nil { + return "", errors.Wrap(err, "converting index to qualified table id") + } + return qtid.Key(), nil +} diff --git a/dax/queryer/compute_importer.go b/dax/queryer/compute_importer.go new file mode 100644 index 000000000..6f22fa8ed --- /dev/null +++ b/dax/queryer/compute_importer.go @@ -0,0 +1,68 @@ +package queryer + +import ( + "context" + + featurebase "github.com/molecula/featurebase/v3" + "github.com/molecula/featurebase/v3/client" + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/errors" +) + +// Ensure type implements interface. +var _ Importer = &ComputeImporter{} + +// ComputeImporter is an implementation of the Importer interface which uses a +// featurebase client to communicate with the compute node. +type ComputeImporter struct { + addr dax.Address +} + +func NewComputeImporter(addr dax.Address) *ComputeImporter { + return &ComputeImporter{ + addr: addr, + } +} + +func (ci *ComputeImporter) CreateIndexKeys(ctx context.Context, index string, keys ...string) (map[string]uint64, error) { + fbClient, err := fbClient(ci.addr) + if err != nil { + return nil, errors.Wrap(err, "getting fb client") + } + + idx := client.NewIndex(index) + return fbClient.CreateIndexKeys(idx, keys...) +} + +func (ci *ComputeImporter) CreateFieldKeys(ctx context.Context, index, field string, keys ...string) (map[string]uint64, error) { + fbClient, err := fbClient(ci.addr) + if err != nil { + return nil, errors.Wrap(err, "getting fb client") + } + + idx := client.NewIndex(index) + fld := idx.Field(field) + return fbClient.CreateFieldKeys(fld, keys...) +} + +func (ci *ComputeImporter) Import(ctx context.Context, req *featurebase.ImportRequest, opts ...featurebase.ImportOption) error { + fbClient, err := fbClient(ci.addr) + if err != nil { + return errors.Wrap(err, "getting fb client") + } + + idx := client.NewIndex(req.Index) + fld := idx.Field(req.Field) + return fbClient.Import(fld, req.Shard, req.RowIDs, req.ColumnIDs, req.Clear) +} + +func (ci *ComputeImporter) ImportValue(ctx context.Context, req *featurebase.ImportValueRequest, opts ...featurebase.ImportOption) error { + fbClient, err := fbClient(ci.addr) + if err != nil { + return errors.Wrap(err, "getting fb client") + } + + idx := client.NewIndex(req.Index) + fld := idx.Field(req.Field) + return fbClient.ImportValues(fld, req.Shard, req.Values, req.ColumnIDs, req.Clear) +} diff --git a/dax/queryer/config.go b/dax/queryer/config.go new file mode 100644 index 000000000..6efb58f49 --- /dev/null +++ b/dax/queryer/config.go @@ -0,0 +1,21 @@ +package queryer + +import ( + "github.com/molecula/featurebase/v3/logger" +) + +// Config defines the configuration parameters for Queryer. At the moment, it's +// being used to serve two different purposes. This first is to provide the +// config parameters for the toml (i.e. human-friendly) file used at server +// startup. The second is as the Config for the Queryer type. If this gets more +// complex, it might make sense to split this into two different config structs. +// We initially did that with something called "Injections", but that separation +// was a bit premature. +type Config struct { + MDSAddress string `toml:"mds-address"` + + MDS MDS `toml:"-"` + Router Router `toml:"-"` + + Logger logger.Logger `toml:"-"` +} diff --git a/dax/queryer/featurebase_importer.go b/dax/queryer/featurebase_importer.go new file mode 100644 index 000000000..087df1a8d --- /dev/null +++ b/dax/queryer/featurebase_importer.go @@ -0,0 +1,41 @@ +package queryer + +import ( + "context" + + featurebase "github.com/molecula/featurebase/v3" +) + +// Ensure type implements interface. +var _ Importer = &FeatureBaseImporter{} + +// FeatureBaseImporter is an implementation of the Importer interface which uses +// a pointer to a featurebase.API to make the underlying calls. This assumes +// those calls need to be Qcx aware, so this takes that into account. +type FeatureBaseImporter struct { + api *featurebase.API +} + +func NewFeatureBaseImporter(api *featurebase.API) *FeatureBaseImporter { + return &FeatureBaseImporter{ + api: api, + } +} + +func (fi *FeatureBaseImporter) CreateIndexKeys(ctx context.Context, index string, keys ...string) (map[string]uint64, error) { + return fi.api.CreateIndexKeys(ctx, index, keys...) +} + +func (fi *FeatureBaseImporter) CreateFieldKeys(ctx context.Context, index, field string, keys ...string) (map[string]uint64, error) { + return fi.api.CreateFieldKeys(ctx, index, field, keys...) +} + +func (fi *FeatureBaseImporter) Import(ctx context.Context, req *featurebase.ImportRequest, opts ...featurebase.ImportOption) error { + qcx := fi.api.Txf().NewQcx() + return fi.api.Import(ctx, qcx, req, opts...) +} + +func (fi *FeatureBaseImporter) ImportValue(ctx context.Context, req *featurebase.ImportValueRequest, opts ...featurebase.ImportOption) error { + qcx := fi.api.Txf().NewQcx() + return fi.api.ImportValue(ctx, qcx, req, opts...) +} diff --git a/dax/queryer/http/handler.go b/dax/queryer/http/handler.go new file mode 100644 index 000000000..579c7d3f9 --- /dev/null +++ b/dax/queryer/http/handler.go @@ -0,0 +1,113 @@ +package http + +import ( + "encoding/json" + "net/http" + + "github.com/gorilla/mux" + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/queryer" +) + +func Handler(q *queryer.Queryer) http.Handler { + svr := &server{ + queryer: q, + } + + router := mux.NewRouter() + router.HandleFunc("/health", svr.getHealth).Methods("GET").Name("GetHealth") + router.HandleFunc("/query", svr.postQuery).Methods("POST").Name("PostQuery") + + // /sql is a subset of /query, added here to provide an easy integration + // with the FeatureBase cli tool. + router.HandleFunc("/sql", svr.postSQL).Methods("POST").Name("PostSQL") + + return router +} + +type server struct { + queryer *queryer.Queryer +} + +// GET /health +func (s *server) getHealth(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) +} + +// POST /query +func (s *server) postQuery(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + req := QueryRequest{} + if err := json.NewDecoder(body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + ctx := r.Context() + + var resp interface{} + var err error + qual := dax.NewTableQualifier(req.OrganizationID, req.DatabaseID) + if req.SQL != "" { + resp, err = s.queryer.QuerySQL(ctx, qual, req.SQL) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + } else { + resp, err = s.queryer.QueryPQL(ctx, qual, req.Table, req.PQL) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + } + + if err := json.NewEncoder(w).Encode(resp); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } +} + +// POST /sql +func (s *server) postSQL(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + req := SQLRequest{} + if err := json.NewDecoder(body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + ctx := r.Context() + + qual := dax.NewTableQualifier(req.OrganizationID, req.DatabaseID) + resp, err := s.queryer.QuerySQL(ctx, qual, req.SQL) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + if err := json.NewEncoder(w).Encode(resp); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } +} + +type QueryRequest struct { + OrganizationID dax.OrganizationID `json:"org-id"` + DatabaseID dax.DatabaseID `json:"db-id"` + Table dax.TableName `json:"table-name"` + PQL string `json:"pql"` + SQL string `json:"sql"` +} + +type SQLRequest struct { + OrganizationID dax.OrganizationID `json:"org-id"` + DatabaseID dax.DatabaseID `json:"db-id"` + SQL string `json:"sql"` +} + +type QueryResponse interface{} diff --git a/dax/queryer/importer.go b/dax/queryer/importer.go new file mode 100644 index 000000000..05b7facbe --- /dev/null +++ b/dax/queryer/importer.go @@ -0,0 +1,129 @@ +package queryer + +import ( + "context" + "strings" + "time" + + featurebase "github.com/molecula/featurebase/v3" + "github.com/molecula/featurebase/v3/batch" + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/mds/schemar" + "github.com/molecula/featurebase/v3/errors" + "github.com/molecula/featurebase/v3/roaring" +) + +// Ensure type implements interface. +var _ batch.Importer = &batchImporter{} + +func newBatchImporter(importer batch.Importer, qual dax.TableQualifier, schemar schemar.Schemar) *batchImporter { + return &batchImporter{ + importer: importer, + qual: qual, + schemar: schemar, + } +} + +// batchImporter is an implementation of the batch.Importer. It is a wrapper +// around idk/mds/Importer that can take index values which are either indexName +// (like "foo") or TableKey (like "tbl__acme__db1__foo123"). This wrapper looks +// at the value to determine if it is a TableKey or not and converts it +// appropriately. It's kind of annoying; we really need to be certain where +// we're expecting indexName vs TableKey. +type batchImporter struct { + importer batch.Importer + qual dax.TableQualifier + schemar schemar.Schemar +} + +func (b *batchImporter) StartTransaction(ctx context.Context, id string, timeout time.Duration, exclusive bool, requestTimeout time.Duration) (*featurebase.Transaction, error) { + return b.importer.StartTransaction(ctx, id, timeout, exclusive, requestTimeout) +} + +func (b *batchImporter) FinishTransaction(ctx context.Context, id string) (*featurebase.Transaction, error) { + return b.importer.FinishTransaction(ctx, id) +} + +func (b *batchImporter) CreateIndexKeys(ctx context.Context, idx *featurebase.IndexInfo, keys ...string) (map[string]uint64, error) { + // Used as an example: + // qual: [acme:db1] + // INSERT INTO foo VALUE (1, 10) + // Currently, the table name coming through IndexInfo right now is the sql + // table name (i.e. foo, not tbl__acme__db1__foo123). For SELECT queries, + // we're currently doing that conversion in the orchestrator.Execute() + // method (which means that we currently only support a single index in SQL + // queries). Therefore, we need to convert idx.Name to a TableKey. + tkey, err := b.indexToQualifiedTableKey(ctx, idx.Name) + if err != nil { + return nil, errors.Wrapf(err, "converting index to qualified table key: %s", idx.Name) + } + idx.Name = string(tkey) + + return b.importer.CreateIndexKeys(ctx, idx, keys...) +} + +func (b *batchImporter) CreateFieldKeys(ctx context.Context, index string, field *featurebase.FieldInfo, keys ...string) (map[string]uint64, error) { + tkey, err := b.indexToQualifiedTableKey(ctx, index) + if err != nil { + return nil, errors.Wrapf(err, "converting index to qualified table key: %s", index) + } + return b.importer.CreateFieldKeys(ctx, string(tkey), field, keys...) +} + +func (b *batchImporter) ImportRoaringBitmap(ctx context.Context, index string, field *featurebase.FieldInfo, shard uint64, views map[string]*roaring.Bitmap, clear bool) error { + tkey, err := b.indexToQualifiedTableKey(ctx, index) + if err != nil { + return errors.Wrapf(err, "converting index to qualified table key: %s", index) + } + return b.importer.ImportRoaringBitmap(ctx, string(tkey), field, shard, views, clear) +} + +func (b *batchImporter) ImportRoaringShard(ctx context.Context, index string, shard uint64, request *featurebase.ImportRoaringShardRequest) error { + tkey, err := b.indexToQualifiedTableKey(ctx, index) + if err != nil { + return errors.Wrapf(err, "converting index to qualified table key: %s", index) + } + return b.importer.ImportRoaringShard(ctx, string(tkey), shard, request) +} + +func (b *batchImporter) EncodeImportValues(ctx context.Context, index string, field *featurebase.FieldInfo, shard uint64, vals []int64, ids []uint64, clear bool) (path string, data []byte, err error) { + tkey, err := b.indexToQualifiedTableKey(ctx, index) + if err != nil { + return "", nil, errors.Wrapf(err, "converting index to qualified table key: %s", index) + } + return b.importer.EncodeImportValues(ctx, string(tkey), field, shard, vals, ids, clear) +} + +func (b *batchImporter) EncodeImport(ctx context.Context, index string, field *featurebase.FieldInfo, shard uint64, vals, ids []uint64, clear bool) (path string, data []byte, err error) { + tkey, err := b.indexToQualifiedTableKey(ctx, index) + if err != nil { + return "", nil, errors.Wrapf(err, "converting index to qualified table key: %s", index) + } + return b.importer.EncodeImport(ctx, string(tkey), field, shard, vals, ids, clear) +} + +func (b *batchImporter) DoImport(ctx context.Context, index string, field *featurebase.FieldInfo, shard uint64, path string, data []byte) error { + tkey, err := b.indexToQualifiedTableKey(ctx, index) + if err != nil { + return errors.Wrapf(err, "converting index to qualified table key: %s", index) + } + return b.importer.DoImport(ctx, string(tkey), field, shard, path, data) +} + +func (b *batchImporter) StatsTiming(name string, value time.Duration, rate float64) { + b.importer.StatsTiming(name, value, rate) +} + +// TODO(tlt): this method was copied from orchestrator.go. Can we centralize +// this logic? +func (b *batchImporter) indexToQualifiedTableKey(ctx context.Context, index string) (dax.TableKey, error) { + if strings.HasPrefix(index, dax.PrefixTable+dax.TableKeyDelimiter) { + return dax.TableKey(index), nil + } + + qtid, err := b.schemar.TableID(ctx, b.qual, dax.TableName(index)) + if err != nil { + return "", errors.Wrap(err, "converting index to qualified table id") + } + return qtid.Key(), nil +} diff --git a/dax/queryer/interfaces.go b/dax/queryer/interfaces.go new file mode 100644 index 000000000..9abdf95cc --- /dev/null +++ b/dax/queryer/interfaces.go @@ -0,0 +1,78 @@ +package queryer + +import ( + "context" + + featurebase "github.com/molecula/featurebase/v3" + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/mds/controller" + "github.com/molecula/featurebase/v3/dax/mds/schemar" +) + +type MDS interface { + // Controller-related methods. + ComputeNodes(ctx context.Context, qtid dax.QualifiedTableID, shards ...dax.ShardNum) ([]controller.ComputeNode, error) + IngestPartition(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum) (dax.Address, error) + IngestShard(ctx context.Context, qtid dax.QualifiedTableID, shard dax.ShardNum) (dax.Address, error) + TranslateNodes(ctx context.Context, qtid dax.QualifiedTableID, partitions ...dax.PartitionNum) ([]controller.TranslateNode, error) + + // Schemar-related methods. + schemar.Schemar +} + +// Ensure type implements interface. +var _ MDS = &NopMDS{} + +// NopMDS is a no-op implementation of the MDS interface. +type NopMDS struct { + schemar.NopSchemar +} + +func NewNopMDS() *NopMDS { + return &NopMDS{} +} + +func (m *NopMDS) ComputeNodes(ctx context.Context, qtid dax.QualifiedTableID, shards ...dax.ShardNum) ([]controller.ComputeNode, error) { + return nil, nil +} +func (m *NopMDS) IngestPartition(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum) (dax.Address, error) { + return "", nil +} +func (m *NopMDS) IngestShard(ctx context.Context, qtid dax.QualifiedTableID, shard dax.ShardNum) (dax.Address, error) { + return "", nil +} +func (m *NopMDS) TranslateNodes(ctx context.Context, qtid dax.QualifiedTableID, partitions ...dax.PartitionNum) ([]controller.TranslateNode, error) { + return nil, nil +} + +type Importer interface { + CreateIndexKeys(ctx context.Context, index string, keys ...string) (map[string]uint64, error) + CreateFieldKeys(ctx context.Context, index, field string, keys ...string) (map[string]uint64, error) + Import(ctx context.Context, req *featurebase.ImportRequest, opts ...featurebase.ImportOption) error + ImportValue(ctx context.Context, req *featurebase.ImportValueRequest, opts ...featurebase.ImportOption) error +} + +// Ensure type implements interface. +var _ Importer = &NopImporter{} + +// NopImporter is a no-op implementation of the Importer interface. +type NopImporter struct{} + +func NewNopImporter() *NopImporter { + return &NopImporter{} +} + +func (n *NopImporter) CreateIndexKeys(ctx context.Context, index string, keys ...string) (map[string]uint64, error) { + return nil, nil +} + +func (n *NopImporter) CreateFieldKeys(ctx context.Context, index, field string, keys ...string) (map[string]uint64, error) { + return nil, nil +} + +func (n *NopImporter) Import(ctx context.Context, req *featurebase.ImportRequest, opts ...featurebase.ImportOption) error { + return nil +} +func (n *NopImporter) ImportValue(ctx context.Context, req *featurebase.ImportValueRequest, opts ...featurebase.ImportOption) error { + return nil +} diff --git a/dax/queryer/orchestrator.go b/dax/queryer/orchestrator.go new file mode 100644 index 000000000..3c7af8eb5 --- /dev/null +++ b/dax/queryer/orchestrator.go @@ -0,0 +1,3494 @@ +// Copyright 2021 Molecula Corp. All rights reserved. +package queryer + +import ( + "context" + "fmt" + "math" + "sort" + "strings" + "time" + + featurebase "github.com/molecula/featurebase/v3" + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/mds/controller" + "github.com/molecula/featurebase/v3/dax/mds/schemar" + "github.com/molecula/featurebase/v3/errors" + "github.com/molecula/featurebase/v3/logger" + "github.com/molecula/featurebase/v3/net" + "github.com/molecula/featurebase/v3/pql" + "github.com/molecula/featurebase/v3/stats" + "github.com/molecula/featurebase/v3/tracing" + "golang.org/x/sync/errgroup" +) + +// Field types. +const ( + FieldTypeSet = "set" + FieldTypeInt = "int" + FieldTypeTime = "time" + FieldTypeMutex = "mutex" + FieldTypeBool = "bool" + FieldTypeDecimal = "decimal" + FieldTypeTimestamp = "timestamp" + + // Row ids used for boolean fields. + falseRowID = uint64(0) + trueRowID = uint64(1) +) + +var ErrFieldNotFound error = dax.NewErrFieldDoesNotExist("") + +const ( + errConnectionRefused = "connect: connection refused" +) + +type Topologer interface { + ComputeNodes(ctx context.Context, index string, shards []uint64) ([]controller.ComputeNode, error) +} + +type MDSTopology struct { + mds MDS +} + +func (m *MDSTopology) ComputeNodes(ctx context.Context, index string, shards []uint64) ([]controller.ComputeNode, error) { + var daxShards = make(dax.ShardNums, len(shards)) + for i, s := range shards { + daxShards[i] = dax.ShardNum(s) + } + + qtid := dax.TableKey(index).QualifiedTableID() + + return m.mds.ComputeNodes(ctx, qtid, daxShards...) +} + +// TODO(jaffee) we need version info in here ASAP. whenever schema or topo +// changes, version gets bumped and nodes know to reject queries +// and update their info from the MDS instead of querying it every +// time. +type Translator interface { + CreateIndexKeys(ctx context.Context, index string, keys []string) (map[string]uint64, error) + CreateFieldKeys(ctx context.Context, index string, field string, keys []string) (map[string]uint64, error) + FindIndexKeys(ctx context.Context, index string, keys []string) (map[string]uint64, error) + FindFieldKeys(ctx context.Context, index, field string, keys []string) (map[string]uint64, error) + // TODO(jaffee) the naming here is a cluster. TranslateIndexIDs takes a list, but TranslateFieldIDs takes a set, both have alternate methods that take the other thing. :facepalm: + TranslateIndexIDs(ctx context.Context, index string, ids []uint64) ([]string, error) + TranslateIndexIDSet(ctx context.Context, index string, ids map[uint64]struct{}) (map[uint64]string, error) + TranslateFieldIDs(ctx context.Context, index, field string, ids map[uint64]struct{}) (map[uint64]string, error) + TranslateFieldListIDs(ctx context.Context, index, field string, ids []uint64) ([]string, error) +} + +// executor recursively executes calls in a PQL query across all shards. +type orchestrator struct { + schema featurebase.SchemaInfoAPI + topology Topologer + trans Translator + + // Client used for remote requests. + client *featurebase.InternalClient + + stats stats.StatsClient + logger logger.Logger +} + +func emptyResult(c *pql.Call) interface{} { + switch c.Name { + case "Clear", "ClearRow": + return false + case "Row": + return &featurebase.Row{Keys: []string{}} + case "Rows": + return featurebase.RowIdentifiers{Keys: []string{}} + case "IncludesColumn": + return false + } + return nil +} + +// Execute executes a PQL query. +func (o *orchestrator) Execute(ctx context.Context, index string, q *pql.Query, shards []uint64, opt *featurebase.ExecOptions) (featurebase.QueryResponse, error) { + span, ctx := tracing.StartSpanFromContext(ctx, "orchestrator.Execute") + span.LogKV("pql", q.String()) + defer span.Finish() + + resp := featurebase.QueryResponse{} + + // Check for query cancellation. + if err := validateQueryContext(ctx); err != nil { + return resp, err + } + + // Verify that an index is set. + if index == "" { + return resp, featurebase.ErrIndexRequired + } + + idx, err := o.schema.IndexInfo(ctx, index) + if err != nil { + return resp, errors.Wrap(err, "getting index") + } + + // Default options. + if opt == nil { + opt = &featurebase.ExecOptions{} + } + + results, err := o.execute(ctx, index, q, shards, opt) + if err != nil { + return resp, err + } else if err := validateQueryContext(ctx); err != nil { + return resp, err + } + resp.Results = results + + if err := o.translateResults(ctx, index, idx, q.Calls, results, opt.MaxMemory); err != nil { + if errors.Cause(err) == featurebase.ErrTranslatingKeyNotFound { + // No error - return empty result + resp.Results = make([]interface{}, len(q.Calls)) + for i, c := range q.Calls { + resp.Results[i] = emptyResult(c) + } + return resp, nil + } + return resp, err + } else if err := validateQueryContext(ctx); err != nil { + return resp, err + } + + return resp, nil +} + +func (o *orchestrator) execute(ctx context.Context, index string, q *pql.Query, shards []uint64, opt *featurebase.ExecOptions) ([]interface{}, error) { + span, ctx := tracing.StartSpanFromContext(ctx, "Executor.execute") + defer span.Finish() + + // Apply translations if necessary. + var colTranslations map[string]map[string]uint64 // colID := colTranslations[index][key] + var rowTranslations map[string]map[string]map[string]uint64 // rowID := rowTranslations[index][field][key] + if !opt.Remote { + cols, rows, err := o.preTranslate(ctx, index, q.Calls...) + if err != nil { + return nil, err + } + colTranslations, rowTranslations = cols, rows + } + + // Execute each call serially. + results := make([]interface{}, 0, len(q.Calls)) + for i, call := range q.Calls { + + if err := validateQueryContext(ctx); err != nil { + return nil, err + } + + // Apply call translation. + if !opt.Remote && !opt.PreTranslated { + translated, err := o.translateCall(ctx, call, index, colTranslations, rowTranslations) + if err != nil { + return nil, errors.Wrap(err, "translating call") + } + if translated == nil { + results = append(results, emptyResult(call)) + continue + } + + call = translated + } + + // If you actually make a top-level Distinct call, you + // want a featurebase.SignedRow back. Otherwise, it's something else + // that will be using it as a row, and we only care + // about the positive values, because only positive values + // are valid column IDs. So we don't actually eat top-level + // pre calls. + if call.Name == "Count" { + // Handle count specially, skipping the level directly underneath it. + for _, child := range call.Children { + err := o.handlePreCallChildren(ctx, index, child, shards, opt) + if err != nil { + return nil, err + } + } + } else { + err := o.handlePreCallChildren(ctx, index, call, shards, opt) + if err != nil { + return nil, err + } + } + var v interface{} + var err error + // Top-level calls don't need to precompute cross-index things, + // because we can just pick whatever index we want, but we + // still need to handle them. Since everything else was + // already precomputed by handlePreCallChildren, though, + // we don't need this logic in executeCall. + newIndex := call.CallIndex() + if newIndex != "" && newIndex != index { + v, err = o.executeCall(ctx, newIndex, call, nil, opt) + } else { + v, err = o.executeCall(ctx, index, call, shards, opt) + } + if err != nil { + return nil, err + } + + if vc, ok := v.(featurebase.ValCount); ok { + vc.Cleanup() + v = vc + } + + results = append(results, v) + // Some Calls can have significant data associated with them + // that gets generated during processing, such as Precomputed + // values. Dumping the precomputed data, if any, lets the GC + // free the memory before we get there. + o.dumpPrecomputedCalls(ctx, q.Calls[i]) + } + return results, nil +} + +// handlePreCalls traverses the call tree looking for calls that need +// precomputed values (e.g. Distinct, UnionRows, ConstRow...). +func (o *orchestrator) handlePreCalls(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) error { + if c.Name == "Precomputed" { + idx := c.Args["valueidx"].(int64) + if idx >= 0 && idx < int64(len(opt.EmbeddedData)) { + row := opt.EmbeddedData[idx] + c.Precomputed = make(map[uint64]interface{}, len(row.Segments)) + for _, segment := range row.Segments { + c.Precomputed[segment.Shard()] = &featurebase.Row{Segments: []featurebase.RowSegment{segment}} + } + } else { + return fmt.Errorf("no precomputed data! index %d, len %d", idx, len(opt.EmbeddedData)) + } + return nil + } + newIndex := c.CallIndex() + // A cross-index query is handled by precall. This is inefficient, + // but we have to do it for now because shards might be different and + // we haven't implemented the local precalls that would be enough + // in some cases. + // + // This makes simple cross-index queries noticably inefficient. + // + // If you're here because of that: We should be using PrecallLocal + // in cases where the call isn't already PrecallGlobal, and + // PrecallLocal should wait until we're running on a specific node + // to do the farming-out of just the sub-queries it has to run + // for its local shards. + // + // As is, we have one node querying every node, then sending out + // all the data to every node, including the data that node already + // has. We could reduce the actual copying around dramatically, + // but only in the cases where local is good enough -- not something + // like Distinct, where you can't predict output shard for a result + // from the shard being queried. + if newIndex != "" && newIndex != index { + c.Type = pql.PrecallGlobal + index = newIndex + // we need to recompute shards, then + shards = nil + } + if err := o.handlePreCallChildren(ctx, index, c, shards, opt); err != nil { + return err + } + // child calls already handled, no precall for this, so we're done + if c.Type == pql.PrecallNone { + return nil + } + // We don't try to handle sub-calls from here. I'm not 100% + // sure that's right, but I think the fact that they're happening + // inside a precomputed call may mean they need different + // handling. In any event, the sub-calls will get handled by + // the executeCall when it gets to them... + + // We set c to look like a normal call, and actually execute it: + c.Type = pql.PrecallNone + // possibly override call index. + v, err := o.executeCall(ctx, index, c, shards, opt) + if err != nil { + return err + } + var row *featurebase.Row + switch r := v.(type) { + case *featurebase.Row: + row = r + case featurebase.SignedRow: + row = r.Pos + default: + return fmt.Errorf("precomputed call %s returned unexpected non-Row data: %T", c.Name, v) + } + if err := ctx.Err(); err != nil { + return err + } + c.Children = []*pql.Call{} + c.Name = "Precomputed" + c.Args = map[string]interface{}{"valueidx": len(opt.EmbeddedData)} + // stash a copy of the full results, which can be forwarded to other + // shards if the query has to go to them + opt.EmbeddedData = append(opt.EmbeddedData, row) + // and stash a copy locally, so local calls can use it + if row != nil { + c.Precomputed = make(map[uint64]interface{}, len(row.Segments)) + for _, segment := range row.Segments { + c.Precomputed[segment.Shard()] = &featurebase.Row{Segments: []featurebase.RowSegment{segment}} + } + } + return nil +} + +// dumpPrecomputedCalls throws away precomputed call data. this is used so we +// can drop any large data associated with a call once we've processed +// the call. +func (o *orchestrator) dumpPrecomputedCalls(ctx context.Context, c *pql.Call) { + for _, call := range c.Children { + o.dumpPrecomputedCalls(ctx, call) + } + c.Precomputed = nil +} + +// handlePreCallChildren handles any pre-calls in the children of a given call. +func (o *orchestrator) handlePreCallChildren(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) error { + for i := range c.Children { + if err := ctx.Err(); err != nil { + return err + } + if err := o.handlePreCalls(ctx, index, c.Children[i], shards, opt); err != nil { + return err + } + } + for key, val := range c.Args { + // Do not precompute GroupBy aggregates + if key == "aggregate" { + continue + } + // Handle Call() operations which exist inside named arguments, too. + if call, ok := val.(*pql.Call); ok { + if err := ctx.Err(); err != nil { + return err + } + if err := o.handlePreCalls(ctx, index, call, shards, opt); err != nil { + return err + } + } + } + return nil +} + +// preprocessQuery expands any calls that need preprocessing. +func (o *orchestrator) preprocessQuery(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (*pql.Call, error) { + switch c.Name { + case "All": + _, hasLimit, err := c.UintArg("limit") + if err != nil { + return nil, err + } + _, hasOffset, err := c.UintArg("offset") + if err != nil { + return nil, err + } + if !hasLimit && !hasOffset { + return c, nil + } + + // Rewrite the All() w/ limit to Limit(All()). + c.Children = []*pql.Call{ + { + Name: "All", + }, + } + c.Name = "Limit" + return c, nil + + default: + // Recurse through child calls. + out := make([]*pql.Call, len(c.Children)) + var changed bool + for i, child := range c.Children { + res, err := o.preprocessQuery(ctx, index, child, shards, opt) + if err != nil { + return nil, err + } + if res != child { + changed = true + } + out[i] = res + } + if changed { + c = c.Clone() + c.Children = out + } + return c, nil + } +} + +// executeCall executes a call. +func (o *orchestrator) executeCall(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (interface{}, error) { + span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeCall") + defer span.Finish() + + if err := validateQueryContext(ctx); err != nil { + return nil, err + } else if err := o.validateCallArgs(c); err != nil { + return nil, errors.Wrap(err, "validating args") + } + indexTag := "index:" + index + metricName := "query_" + strings.ToLower(c.Name) + "_total" + statFn := func() { + if !opt.Remote { + o.stats.CountWithCustomTags(metricName, 1, 1.0, []string{indexTag}) + } + } + + // Preprocess the query. + c, err := o.preprocessQuery(ctx, index, c, shards, opt) + if err != nil { + return nil, err + } + + switch c.Name { + case "Sum": + statFn() + res, err := o.executeSum(ctx, index, c, shards, opt) + return res, errors.Wrap(err, "executeSum") + case "Min": + statFn() + res, err := o.executeMin(ctx, index, c, shards, opt) + return res, errors.Wrap(err, "executeMin") + case "Max": + statFn() + res, err := o.executeMax(ctx, index, c, shards, opt) + return res, errors.Wrap(err, "executeMax") + case "MinRow": + statFn() + res, err := o.executeMinRow(ctx, index, c, shards, opt) + return res, errors.Wrap(err, "executeMinRow") + case "MaxRow": + statFn() + res, err := o.executeMaxRow(ctx, index, c, shards, opt) + return res, errors.Wrap(err, "executeMaxRow") + // case "Clear": + // statFn() + // res, err := o.executeClearBit(ctx, index, c, opt) + // return res, errors.Wrap(err, "executeClearBit") + // case "ClearRow": + // statFn() + // res, err := o.executeClearRow(ctx, index, c, shards, opt) + // return res, errors.Wrap(err, "executeClearRow") + case "Distinct": + statFn() + res, err := o.executeDistinct(ctx, index, c, shards, opt) + return res, errors.Wrap(err, "executeDistinct") + // case "Store": + // statFn() + // res, err := o.executeSetRow(ctx, index, c, shards, opt) + // return res, errors.Wrap(err, "executeSetRow") + case "Count": + statFn() + res, err := o.executeCount(ctx, index, c, shards, opt) + return res, errors.Wrap(err, "executeCount") + // case "Set": + // statFn() + // res, err := o.executeSet(ctx, index, c, opt) + // return res, errors.Wrap(err, "executeSet") + case "TopK": + statFn() + res, err := o.executeTopK(ctx, index, c, shards, opt) + return res, errors.Wrap(err, "executeTopK") + case "TopN": + statFn() + res, err := o.executeTopN(ctx, index, c, shards, opt) + return res, errors.Wrap(err, "executeTopN") + case "Rows": + statFn() + res, err := o.executeRows(ctx, index, c, shards, opt) + return res, errors.Wrap(err, "executeRows") + case "Extract": + statFn() + res, err := o.executeExtract(ctx, index, c, shards, opt) + return res, errors.Wrap(err, "executeExtract") + case "GroupBy": + statFn() + res, err := o.executeGroupBy(ctx, index, c, shards, opt) + return res, errors.Wrap(err, "executeGroupBy") + case "Options": + statFn() + res, err := o.executeOptionsCall(ctx, index, c, shards, opt) + return res, errors.Wrap(err, "executeOptionsCall") + case "IncludesColumn": + res, err := o.executeIncludesColumnCall(ctx, index, c, shards, opt) + return res, errors.Wrap(err, "executeIncludesColumnCall") + case "FieldValue": + statFn() + res, err := o.executeFieldValueCall(ctx, index, c, shards, opt) + return res, errors.Wrap(err, "executeFieldValueCall") + case "Precomputed": + res, err := o.executePrecomputedCall(ctx, index, c, shards, opt) + return res, errors.Wrap(err, "executePrecomputedCall") + case "UnionRows": + res, err := o.executeUnionRows(ctx, index, c, shards, opt) + return res, errors.Wrap(err, "executeUnionRows") + case "ConstRow": + res, err := o.executeConstRow(ctx, index, c) + return res, errors.Wrap(err, "executeConstRow") + case "Limit": + res, err := o.executeLimitCall(ctx, index, c, shards, opt) + return res, errors.Wrap(err, "executeLimitCall") + case "Percentile": + res, err := o.executePercentile(ctx, index, c, shards, opt) + return res, errors.Wrap(err, "executePercentile") + // case "Delete": + // statFn() //TODO(twg) need this? + // res, err := o.executeDeleteRecords(ctx, index, c, shards, opt) + // return res, errors.Wrap(err, "executeDelete") + default: // o.g. "Row", "Union", "Intersect" or anything that returns a bitmap. + statFn() + res, err := o.executeBitmapCall(ctx, index, c, shards, opt) + return res, errors.Wrap(err, "executeBitmapCall") + } +} + +// validateCallArgs ensures that the value types in call.Args are expected. +func (o *orchestrator) validateCallArgs(c *pql.Call) error { + if _, ok := c.Args["ids"]; ok { + switch v := c.Args["ids"].(type) { + case []int64, []uint64: + // noop + case []interface{}: + b := make([]int64, len(v)) + for i := range v { + b[i] = v[i].(int64) + } + c.Args["ids"] = b + default: + return fmt.Errorf("invalid call.Args[ids]: %s", v) + } + } + return nil +} + +func (o *orchestrator) executeOptionsCall(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (interface{}, error) { + span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeOptionsCall") + defer span.Finish() + + optCopy := &featurebase.ExecOptions{} + *optCopy = *opt + if arg, ok := c.Args["shards"]; ok { + if optShards, ok := arg.([]interface{}); ok { + shards = []uint64{} + for _, s := range optShards { + if shard, ok := s.(int64); ok { + shards = append(shards, uint64(shard)) + } else { + return nil, errors.New(errors.ErrUncoded, "Query(): shards must be a list of unsigned integers") + } + + } + } else { + return nil, errors.New(errors.ErrUncoded, "Query(): shards must be a list of unsigned integers") + } + } + return o.executeCall(ctx, index, c.Children[0], shards, optCopy) +} + +// executeIncludesColumnCall executes an IncludesColumn() call. +func (o *orchestrator) executeIncludesColumnCall(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (bool, error) { + // Get the shard containing the column, since that's the only + // shard that needs to execute this query. + var shard uint64 + col, ok, err := c.UintArg("column") + if err != nil { + return false, errors.Wrap(err, "getting column from args") + } else if !ok { + return false, errors.New(errors.ErrUncoded, "IncludesColumn call must specify a column") + } + shard = col / featurebase.ShardWidth + + // Merge returned results at coordinating node. + reduceFn := func(ctx context.Context, prev, v interface{}) interface{} { + other, _ := prev.(bool) + return other || v.(bool) + } + + result, err := o.mapReduce(ctx, index, []uint64{shard}, c, opt, reduceFn) + if err != nil { + return false, err + } + return result.(bool), nil +} + +// executeFieldValueCall executes a FieldValue() call. +func (o *orchestrator) executeFieldValueCall(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (_ featurebase.ValCount, err error) { + fieldName, ok := c.Args["field"].(string) + if !ok || fieldName == "" { + return featurebase.ValCount{}, featurebase.ErrFieldRequired + } + + colKey, ok := c.Args["column"] + if !ok || colKey == "" { + return featurebase.ValCount{}, featurebase.ErrColumnRequired + } + + colID, ok, err := c.UintArg("column") + if !ok || err != nil { + return featurebase.ValCount{}, errors.Wrap(err, "getting column argument") + } + + shard := colID / featurebase.ShardWidth + + // Select single returned result at coordinating node. + reduceFn := func(ctx context.Context, prev, v interface{}) interface{} { + other, _ := prev.(featurebase.ValCount) + if other.Count == 1 { + return other + } + return v + } + + result, err := o.mapReduce(ctx, index, []uint64{shard}, c, opt, reduceFn) + if err != nil { + return featurebase.ValCount{}, errors.Wrap(err, "map reduce") + } + other, _ := result.(featurebase.ValCount) + + return other, nil +} + +// executeLimitCall executes a Limit() call. +func (o *orchestrator) executeLimitCall(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (*featurebase.Row, error) { + bitmapCall := c.Children[0] + + limit, hasLimit, err := c.UintArg("limit") + if err != nil { + return nil, errors.Wrap(err, "getting limit") + } + offset, _, err := c.UintArg("offset") + if err != nil { + return nil, errors.Wrap(err, "getting offset") + } + + if !hasLimit { + limit = math.MaxUint64 + } + + // Execute bitmap call, storing the full result on this node. + res, err := o.executeCall(ctx, index, bitmapCall, shards, opt) + if err != nil { + return nil, errors.Wrap(err, "limit map reduce") + } + if res == nil { + res = featurebase.NewRow() + } + + result, ok := res.(*featurebase.Row) + if !ok { + return nil, errors.Errorf("expected Row but got %T", result) + } + + if offset != 0 { + i := 0 + var leadingBits []uint64 + for i < len(result.Segments) && offset > 0 { + seg := result.Segments[i] + count := seg.Count() + if count > offset { + data := seg.Columns() + data = data[offset:] + leadingBits = data + i++ + break + } + + offset -= count + i++ + } + row := featurebase.NewRow(leadingBits...) + row.Merge(&featurebase.Row{Segments: result.Segments[i:]}) + result = row + } + if limit < result.Count() { + i := 0 + var trailingBits []uint64 + for i < len(result.Segments) && limit > 0 { + seg := result.Segments[i] + count := seg.Count() + if count > limit { + data := seg.Columns() + data = data[:limit] + trailingBits = data + break + } + + limit -= count + i++ + } + row := featurebase.NewRow(trailingBits...) + row.Merge(&featurebase.Row{Segments: result.Segments[:i]}) + result = row + } + + return result, nil +} + +// executeSum executes a Sum() call. +func (o *orchestrator) executeSum(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (_ featurebase.ValCount, err error) { + span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeSum") + defer span.Finish() + + fieldName, err := c.FirstStringArg("field", "_field") + if err != nil { + return featurebase.ValCount{}, errors.Wrap(err, "Sum(): field required") + } + + if len(c.Children) > 1 { + return featurebase.ValCount{}, errors.New(errors.ErrUncoded, "Sum() only accepts a single bitmap input") + } + + // Merge returned results at coordinating node. + reduceFn := func(ctx context.Context, prev, v interface{}) interface{} { + other, _ := prev.(featurebase.ValCount) + return other.Add(v.(featurebase.ValCount)) + } + + result, err := o.mapReduce(ctx, index, shards, c, opt, reduceFn) + if err != nil { + return featurebase.ValCount{}, err + } + other, _ := result.(featurebase.ValCount) + + if other.Count == 0 { + return featurebase.ValCount{}, nil + } + + // scale summed response if it's a decimal field and this is + // not a remote query (we're about to return to original client). + if !opt.Remote { + field, err := o.schema.FieldInfo(ctx, index, fieldName) + if field == nil { + return featurebase.ValCount{}, errors.Wrapf(err, "%q", fieldName) + } + if field.Options.Type == FieldTypeDecimal { + dec := pql.NewDecimal(other.Val, field.Options.Scale) + other.DecimalVal = &dec + other.FloatVal = 0 + other.Val = 0 + } + } + + return other, nil +} + +// executeDistinct executes a Distinct call on a field. It returns a +// SignedRow for int fields and a *Row for set/mutex/time fields. +func (o *orchestrator) executeDistinct(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (interface{}, error) { + span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeDistinct") + defer span.Finish() + + field, hasField, err := c.StringArg("field") + if err != nil { + return featurebase.SignedRow{}, errors.Wrap(err, "loading field option in Distinct query") + } else if !hasField { + return featurebase.SignedRow{}, fmt.Errorf("missing field option in Distinct query") + } + + // Merge returned results at coordinating node. + reduceFn := func(ctx context.Context, prev, v interface{}) interface{} { + if err := ctx.Err(); err != nil { + return err + } + switch other := prev.(type) { + case featurebase.SignedRow: + return other.Union(v.(featurebase.SignedRow)) + case *featurebase.Row: + if other == nil { + return v + } else if v.(*featurebase.Row) == nil { + return other + } + return other.Union(v.(*featurebase.Row)) + case nil: + return v + case featurebase.DistinctTimestamp: + return other.Union(v.(featurebase.DistinctTimestamp)) + default: + return errors.Errorf("unexpected return type from executeDistinctShard: %+v %T", other, other) + } + } + + result, err := o.mapReduce(ctx, index, shards, c, opt, reduceFn) + if err != nil { + return nil, errors.Wrap(err, "mapReduce") + } + + if other, ok := result.(featurebase.SignedRow); ok { + other.Field = field + } + return result, nil +} + +// executeMin executes a Min() call. +func (o *orchestrator) executeMin(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (_ featurebase.ValCount, err error) { + span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeMin") + defer span.Finish() + + if _, err := c.FirstStringArg("field", "_field"); err != nil { + return featurebase.ValCount{}, errors.Wrap(err, "Min(): field required") + } + + if len(c.Children) > 1 { + return featurebase.ValCount{}, errors.New(errors.ErrUncoded, "Min() only accepts a single bitmap input") + } + + // Merge returned results at coordinating node. + reduceFn := func(ctx context.Context, prev, v interface{}) interface{} { + other, _ := prev.(featurebase.ValCount) + return other.Smaller(v.(featurebase.ValCount)) + } + + result, err := o.mapReduce(ctx, index, shards, c, opt, reduceFn) + if err != nil { + return featurebase.ValCount{}, err + } + other, _ := result.(featurebase.ValCount) + + if other.Count == 0 { + return featurebase.ValCount{}, nil + } + return other, nil +} + +// executeMax executes a Max() call. +func (o *orchestrator) executeMax(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (_ featurebase.ValCount, err error) { + span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeMax") + defer span.Finish() + + if _, err := c.FirstStringArg("field", "_field"); err != nil { + return featurebase.ValCount{}, errors.Wrap(err, "Max(): field required") + } + + if len(c.Children) > 1 { + return featurebase.ValCount{}, errors.New(errors.ErrUncoded, "Max() only accepts a single bitmap input") + } + + // Merge returned results at coordinating node. + reduceFn := func(ctx context.Context, prev, v interface{}) interface{} { + other, _ := prev.(featurebase.ValCount) + return other.Larger(v.(featurebase.ValCount)) + } + + result, err := o.mapReduce(ctx, index, shards, c, opt, reduceFn) + if err != nil { + return featurebase.ValCount{}, err + } + other, _ := result.(featurebase.ValCount) + + if other.Count == 0 { + return featurebase.ValCount{}, nil + } + return other, nil +} + +// TODO(jaffee) fix this... valcountize assumes access to field details like base +// executePercentile executes a Percentile() call. +func (o *orchestrator) executePercentile(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (_ featurebase.ValCount, err error) { + span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executePercentile") + defer span.Finish() + + // get nth + var nthFloat float64 + nthArg, ok := c.Args["nth"] + if !ok { + return featurebase.ValCount{}, errors.New(errors.ErrUncoded, "Percentile(): nth required") + } + switch nthArg := nthArg.(type) { + case pql.Decimal: + nthFloat = nthArg.Float64() + case int64: + nthFloat = float64(nthArg) + default: + return featurebase.ValCount{}, errors.Errorf("Percentile(): invalid nth='%v' of type (%[1]T), should be a number between 0 and 100 inclusive", c.Args["nth"]) + } + if nthFloat < 0 || nthFloat > 100.0 { + return featurebase.ValCount{}, errors.Errorf("Percentile(): invalid nth value (%f), should be a number between 0 and 100 inclusive", nthFloat) + } + + // get field + fieldName, err := c.FirstStringArg("field", "_field") + if err != nil { + return featurebase.ValCount{}, errors.New(errors.ErrUncoded, "Percentile(): field required") + } + field, err := o.schema.FieldInfo(ctx, index, fieldName) + if err != nil { + return featurebase.ValCount{}, ErrFieldNotFound + } + + // filter call for min & max + var filterCall *pql.Call + + // check if filter provided + if filterArg, ok := c.Args["filter"].(*pql.Call); ok && filterArg != nil { + filterCall = filterArg + } + + // get min + q, _ := pql.ParseString(fmt.Sprintf(`Min(field="%s")`, fieldName)) + minCall := q.Calls[0] + if filterCall != nil { + minCall.Children = append(minCall.Children, filterCall) + } + minVal, err := o.executeMin(ctx, index, minCall, shards, opt) + if err != nil { + return featurebase.ValCount{}, errors.Wrap(err, "executing Min call for Percentile") + } + if nthFloat == 0.0 { + return minVal, nil + } + + // get max + q, _ = pql.ParseString(fmt.Sprintf(`Max(field="%s")`, fieldName)) + maxCall := q.Calls[0] + if filterCall != nil { + maxCall.Children = append(maxCall.Children, filterCall) + } + maxVal, err := o.executeMax(ctx, index, maxCall, shards, opt) + if err != nil { + return featurebase.ValCount{}, errors.Wrap(err, "executing Max call for Percentile") + } + // set up reusables + var countCall, rangeCall *pql.Call + if filterCall == nil { + countQuery, _ := pql.ParseString(fmt.Sprintf("Count(Row(%s < 0))", fieldName)) + countCall = countQuery.Calls[0] + rangeCall = countCall.Children[0] + } else { + countQuery, _ := pql.ParseString(fmt.Sprintf(`Count(Intersect(Row(%s < 0)))`, fieldName)) + countCall = countQuery.Calls[0] + intersectCall := countCall.Children[0] + intersectCall.Children = append(intersectCall.Children, filterCall) + rangeCall = intersectCall.Children[0] + } + + k := (100 - nthFloat) / nthFloat + + min, max := minVal.Val, maxVal.Val + // estimate nth val, eg median when nth=0.5 + for min < max { + // compute average without integer overflow, then correct for division of + // odd numbers by 2 + possibleNthVal := ((max / 2) + (min / 2)) + (((max % 2) + (min % 2)) / 2) + // possibleNthVal = (max + min) / 2 + // get left count + rangeCall.Args[fieldName] = &pql.Condition{ + Op: pql.Token(pql.LT), + Value: possibleNthVal, + } + leftCountUint64, err := o.executeCount(ctx, index, countCall, shards, opt) + if err != nil { + return featurebase.ValCount{}, errors.Wrap(err, "executing Count call L for Percentile") + } + leftCount := int64(leftCountUint64) + + // get right count + rangeCall.Args[fieldName] = &pql.Condition{ + Op: pql.Token(pql.GT), + Value: possibleNthVal, + } + rightCountUint64, err := o.executeCount(ctx, index, countCall, shards, opt) + if err != nil { + return featurebase.ValCount{}, errors.Wrap(err, "executing Count call R for Percentile") + } + rightCount := int64(rightCountUint64) + + // 'weight' the left count as per k + leftCountWeighted := int64(math.Round(k * float64(leftCount))) + + // binary search + if leftCountWeighted > rightCount { + max = possibleNthVal - 1 + } else if leftCountWeighted < rightCount { + min = possibleNthVal + 1 + } else { + return cookValCount(possibleNthVal, 1, field), nil + } + } + + return cookValCount(min, 1, field), nil +} + +func cookValCount(val int64, cnt uint64, field *featurebase.FieldInfo) featurebase.ValCount { + valCount := featurebase.ValCount{Count: int64(cnt)} + base := field.Options.Base + switch field.Options.Type { + case featurebase.FieldTypeDecimal: + dec := pql.NewDecimal(val+base, field.Options.Scale) + valCount.DecimalVal = &dec + case FieldTypeTimestamp: + valCount.TimestampVal = time.Unix(0, (val+base)*featurebase.TimeUnitNanos(field.Options.TimeUnit)).UTC() + } + valCount.Val = val + base + return valCount +} + +// executeMinRow executes a MinRow() call. +func (o *orchestrator) executeMinRow(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (_ interface{}, err error) { + span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeMinRow") + defer span.Finish() + + if field := c.Args["field"]; field == "" { + return featurebase.ValCount{}, errors.New(errors.ErrUncoded, "MinRow(): field required") + } + + // Merge returned results at coordinating node. + reduceFn := func(ctx context.Context, prev, v interface{}) interface{} { + // if minRowID exists, and if it is smaller than the other one return it. + // otherwise return the minRowID of the one which exists. + if prev == nil { + return v + } else if v == nil { + return prev + } + prevp, _ := prev.(featurebase.PairField) + vp, _ := v.(featurebase.PairField) + if prevp.Pair.Count > 0 && vp.Pair.Count > 0 { + if prevp.Pair.ID < vp.Pair.ID { + return prevp + } + return vp + } else if prevp.Pair.Count > 0 { + return prevp + } + return vp + } + + return o.mapReduce(ctx, index, shards, c, opt, reduceFn) +} + +// executeMaxRow executes a MaxRow() call. +func (o *orchestrator) executeMaxRow(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (_ interface{}, err error) { + span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeMaxRow") + defer span.Finish() + + if field := c.Args["field"]; field == "" { + return featurebase.ValCount{}, errors.New(errors.ErrUncoded, "MaxRow(): field required") + } + + // Merge returned results at coordinating node. + reduceFn := func(ctx context.Context, prev, v interface{}) interface{} { + // if minRowID exists, and if it is smaller than the other one return it. + // otherwise return the minRowID of the one which exists. + if prev == nil { + return v + } else if v == nil { + return prev + } + prevp, _ := prev.(featurebase.PairField) + vp, _ := v.(featurebase.PairField) + if prevp.Pair.Count > 0 && vp.Pair.Count > 0 { + if prevp.Pair.ID > vp.Pair.ID { + return prevp + } + return vp + } else if prevp.Pair.Count > 0 { + return prevp + } + return vp + } + + return o.mapReduce(ctx, index, shards, c, opt, reduceFn) +} + +// executePrecomputedCall pretends to execute a call that we have a precomputed value for. +func (o *orchestrator) executePrecomputedCall(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (_ *featurebase.Row, err error) { + span, _ := tracing.StartSpanFromContext(ctx, "Executor.executePrecomputedCall") + defer span.Finish() + result := featurebase.NewRow() + + for _, row := range c.Precomputed { + result.Merge(row.(*featurebase.Row)) + } + return result, nil +} + +// executeBitmapCall executes a call that returns a bitmap. +func (o *orchestrator) executeBitmapCall(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (_ *featurebase.Row, err error) { + + span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeBitmapCall") + span.LogKV("pqlCallName", c.Name) + defer span.Finish() + + indexTag := "index:" + index + metricName := "query_" + strings.ToLower(c.Name) + "_total" + if c.Name == "Row" && c.HasConditionArg() { + metricName = "query_row_bsi_total" + } + if !opt.Remote { + o.stats.CountWithCustomTags(metricName, 1, 1.0, []string{indexTag}) + } + + // Merge returned results at coordinating node. + reduceFn := func(ctx context.Context, prev, v interface{}) interface{} { + other, _ := prev.(*featurebase.Row) + if other == nil { + // TODO... what's going on on the following line + other = featurebase.NewRow() // bug! this row ends up containing Badger Txn data that should be accessed outside the Txn. + } + if err := ctx.Err(); err != nil { + return err + } + other.Merge(v.(*featurebase.Row)) + return other + } + + other, err := o.mapReduce(ctx, index, shards, c, opt, reduceFn) + if err != nil { + return nil, errors.Wrap(err, "map reduce") + } + + row, _ := other.(*featurebase.Row) + + return row, nil +} + +type Error string // TODO(jaffee) convert to standard error package + +func (e Error) Error() string { return string(e) } + +const ViewNotFound = Error("view not found") +const FragmentNotFound = Error("fragment not found") + +func (o *orchestrator) executeTopK(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (interface{}, error) { + span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeTopK") + defer span.Finish() + + reduceFn := func(ctx context.Context, prev, v interface{}) interface{} { + x, _ := prev.([]*featurebase.Row) + y, _ := v.([]*featurebase.Row) + return ([]*featurebase.Row)(featurebase.AddBSI(x, y)) + } + + other, err := o.mapReduce(ctx, index, shards, c, opt, reduceFn) + if err != nil { + return nil, err + } + results, _ := other.([]*featurebase.Row) + + if opt.Remote { + return results, nil + } + + k, hasK, err := c.UintArg("k") + if err != nil { + return nil, errors.Wrap(err, "fetching k") + } + + var limit *uint64 + if hasK { + limit = &k + } + + var dst []featurebase.Pair + featurebase.BSIData(results).PivotDescending(featurebase.NewRow().Union(results...), 0, limit, nil, func(count uint64, ids ...uint64) { + for _, id := range ids { + dst = append(dst, featurebase.Pair{ + ID: id, + Count: count, + }) + } + }) + + fieldName, hasFieldName, err := c.StringArg("_field") + if err != nil { + return nil, errors.Wrap(err, "fetching TopK field") + } else if !hasFieldName { + return nil, errors.New(errors.ErrUncoded, "missing field in TopK") + } + + return &featurebase.PairsField{ + Pairs: dst, + Field: fieldName, + }, nil +} + +// uint64Slice represents a sortable slice of uint64 numbers. +type uint64Slice []uint64 + +func (p uint64Slice) Swap(i, j int) { p[i], p[j] = p[j], p[i] } +func (p uint64Slice) Len() int { return len(p) } +func (p uint64Slice) Less(i, j int) bool { return p[i] < p[j] } + +// executeTopN executes a TopN() call. +// This first performs the TopN() to determine the top results and then +// requeries to retrieve the full counts for each of the top results. +func (o *orchestrator) executeTopN(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (*featurebase.PairsField, error) { + span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeTopN") + defer span.Finish() + + idsArg, _, err := c.UintSliceArg("ids") + if err != nil { + return nil, fmt.Errorf("executeTopN: %v", err) + } + + fieldName, _ := c.Args["_field"].(string) + n, _, err := c.UintArg("n") + if err != nil { + return nil, fmt.Errorf("executeTopN: %v", err) + } + + // Execute original query. + pairs, err := o.executeTopNShards(ctx, index, c, shards, opt) + if err != nil { + return nil, errors.Wrap(err, "finding top results") + } + + // If this call is against specific ids, or we didn't get results, + // or we are part of a larger distributed query then don't refetch. + if len(pairs.Pairs) == 0 || len(idsArg) > 0 || opt.Remote { + return &featurebase.PairsField{ + Pairs: pairs.Pairs, + Field: fieldName, + }, nil + } + // Only the original caller should refetch the full counts. + // TODO(@kuba--): ...but do we really need `Clone` here? + other := c.Clone() + + ids := featurebase.Pairs(pairs.Pairs).Keys() + sort.Sort(uint64Slice(ids)) + other.Args["ids"] = ids + + trimmedList, err := o.executeTopNShards(ctx, index, other, shards, opt) + if err != nil { + return nil, errors.Wrap(err, "retrieving full counts") + } + + if n != 0 && int(n) < len(trimmedList.Pairs) { + trimmedList.Pairs = trimmedList.Pairs[0:n] + } + + return &featurebase.PairsField{ + Pairs: trimmedList.Pairs, + Field: fieldName, + }, nil +} + +func (o *orchestrator) executeTopNShards(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (*featurebase.PairsField, error) { + span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeTopNShards") + defer span.Finish() + + // Merge returned results at coordinating node. + reduceFn := func(ctx context.Context, prev, v interface{}) interface{} { + other, _ := prev.(*featurebase.PairsField) + vpf, _ := v.(*featurebase.PairsField) + if other == nil { + return vpf + } else if vpf == nil { + return other + } + if err := ctx.Err(); err != nil { + return err + } + other.Pairs = featurebase.Pairs(other.Pairs).Add(vpf.Pairs) + return other + } + + other, err := o.mapReduce(ctx, index, shards, c, opt, reduceFn) + if err != nil { + return nil, err + } + results, _ := other.(*featurebase.PairsField) + + // Sort final merged results. + sort.Sort(featurebase.Pairs(results.Pairs)) + + return results, nil +} + +// order denotes sort order—can be asc or desc (see constants below). +type order bool + +const ( + asc order = true + desc order = false +) + +// groupCountSorter sorts the output of a GroupBy request (a +// []GroupCount) according to sorting instructions encoded in "fields" +// and "order". +// +// Each field in "fields" is an integer which can be -1 to denote +// sorting on the Count and -2 to denote sorting on the +// sum/aggregate. Currently nothing else is supported, but the idea +// was that if there were positive integers they would be indexes into +// GroupCount.FieldRow and allowing sorting on the values of different +// fields in the group. Each item in "order" corresponds to the same +// index in "fields" and denotes the order of the sort. +type groupCountSorter struct { + fields []int + order []order + data []featurebase.GroupCount +} + +func (g *groupCountSorter) Len() int { return len(g.data) } +func (g *groupCountSorter) Swap(i, j int) { g.data[i], g.data[j] = g.data[j], g.data[i] } +func (g *groupCountSorter) Less(i, j int) bool { + gci, gcj := g.data[i], g.data[j] + for idx, fieldIndex := range g.fields { + fieldOrder := g.order[idx] + switch fieldIndex { + case -1: // Count + if gci.Count < gcj.Count { + return fieldOrder == asc + } else if gci.Count > gcj.Count { + return fieldOrder == desc + } + case -2: // Aggregate + if gci.Agg < gcj.Agg { + return fieldOrder == asc + } else if gci.Agg > gcj.Agg { + return fieldOrder == desc + } + default: + panic("impossible") + } + } + return false +} + +// getSorter hackily parses the sortSpec and figures out how to sort +// the GroupBy results. +func getSorter(sortSpec string) (*groupCountSorter, error) { + gcs := &groupCountSorter{ + fields: []int{}, + order: []order{}, + } + sortOn := strings.Split(sortSpec, ",") + for _, sortField := range sortOn { + sortField = strings.TrimSpace(sortField) + fieldDir := strings.Fields(sortField) + if len(fieldDir) == 0 { + return nil, errors.Errorf("invalid sorting directive: '%s'", sortField) + } else if fieldDir[0] == "count" { + gcs.fields = append(gcs.fields, -1) + } else if fieldDir[0] == "aggregate" || fieldDir[0] == "sum" { + gcs.fields = append(gcs.fields, -2) + } else { + return nil, errors.Errorf("sorting is only supported on count, aggregate, or sum, not '%s'", fieldDir[0]) + } + + if len(fieldDir) == 1 { + gcs.order = append(gcs.order, desc) + } else if len(fieldDir) > 2 { + return nil, errors.Errorf("parsing sort directive: '%s': too many elements", sortField) + } else if fieldDir[1] == "asc" { + gcs.order = append(gcs.order, asc) + } else if fieldDir[1] == "desc" { + gcs.order = append(gcs.order, desc) + } else { + return nil, errors.Errorf("unknown sort direction '%s'", fieldDir[1]) + } + } + return gcs, nil +} + +// findGroupCounts gets a safe-to-use but possibly empty []GroupCount from +// an interface which might be a *GroupCounts or a []GroupCount. +func findGroupCounts(v interface{}) []featurebase.GroupCount { + switch gc := v.(type) { + case []featurebase.GroupCount: + return gc + case *featurebase.GroupCounts: + return gc.Groups() + } + return nil +} + +func (o *orchestrator) executeGroupBy(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (*featurebase.GroupCounts, error) { + span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeGroupBy") + defer span.Finish() + // validate call + if len(c.Children) == 0 { + return nil, errors.New(errors.ErrUncoded, "need at least one child call") + } + limit := int(^uint(0) >> 1) + if lim, hasLimit, err := c.UintArg("limit"); err != nil { + return nil, err + } else if hasLimit { + limit = int(lim) + } + filter, _, err := c.CallArg("filter") + if err != nil { + return nil, err + } + + var sorter *groupCountSorter + if sortSpec, found, err := c.StringArg("sort"); err != nil { + return nil, errors.Wrap(err, "getting sort arg") + } else if found { + sorter, err = getSorter(sortSpec) + if err != nil { + return nil, errors.Wrap(err, "parsing sort spec") + } + // don't want to prematurely limit the results if we're sorting + limit = int(^uint(0) >> 1) + } + having, hasHaving, err := c.CallArg("having") + if err != nil { + return nil, errors.Wrap(err, "getting 'having' argument") + } else if hasHaving { + // don't want to prematurely limit the results if we're filtering some out + limit = int(^uint(0) >> 1) + } + + // perform necessary Rows queries (any that have limit or columns args) - + // TODO, call async? would only help if multiple Rows queries had a column + // or limit arg. + // TODO support TopN in here would be really cool - and pretty easy I think. + childRows := make([]featurebase.RowIDs, len(c.Children)) + for i, child := range c.Children { + // Check "field" first for backwards compatibility, then set _field. + // TODO: remove at Pilosa 2.0 + if fieldName, ok := child.Args["field"].(string); ok { + child.Args["_field"] = fieldName + } + + if child.Name != "Rows" { + return nil, errors.Errorf("'%s' is not a valid child query for GroupBy, must be 'Rows'", child.Name) + } + _, hasLimit, err := child.UintArg("limit") + if err != nil { + return nil, errors.Wrap(err, "getting limit") + } + _, hasCol, err := child.UintArg("column") + if err != nil { + return nil, errors.Wrap(err, "getting column") + } + _, hasLike, err := child.StringArg("like") + if err != nil { + return nil, errors.Wrap(err, "getting like") + } + _, hasIn, err := child.UintSliceArg("in") + if err != nil { + return nil, errors.Wrap(err, "getting 'in'") + } + + if hasLimit || hasCol || hasLike || hasIn { // we need to perform this query cluster-wide ahead of executeGroupByShard + if idx, ok := child.Args["valueidx"].(int64); ok { + // The rows query was already completed on the initiating node. + childRows[i] = opt.EmbeddedData[idx].Columns() + continue + } + + r, er := o.executeRows(ctx, index, child, shards, opt) + if er != nil { + return nil, errors.Wrap(er, "getting rows for ") + } + // need to sort because filters assume ordering + sort.Slice(r, func(x, y int) bool { return r[x] < r[y] }) + childRows[i] = r + if len(childRows[i]) == 0 { // there are no results because this field has no values. + return &featurebase.GroupCounts{}, nil + } + + // Stuff the result into opt.EmbeddedData so that it gets sent to other nodes in the map-reduce. + // This is flagged as "NoSplit" to ensure that the entire row gets sent out. + rowsRow := featurebase.NewRow(childRows[i]...) + rowsRow.NoSplit = true + child.Args["valueidx"] = int64(len(opt.EmbeddedData)) + opt.EmbeddedData = append(opt.EmbeddedData, rowsRow) + } + } + + // Merge returned results at coordinating node. + reduceFn := func(ctx context.Context, prev, v interface{}) interface{} { + other := findGroupCounts(prev) + if err := ctx.Err(); err != nil { + return err + } + return mergeGroupCounts(other, findGroupCounts(v), limit) + } + // Get full result set. + other, err := o.mapReduce(ctx, index, shards, c, opt, reduceFn) + if err != nil { + return nil, errors.Wrap(err, "mapReduce") + } + results, _ := other.([]featurebase.GroupCount) + + // If there's no sorting, we want to apply limits before + // calculating the Distinct aggregate which is expensive on a + // per-result basis. + if sorter == nil && !hasHaving { + results, err = applyLimitAndOffsetToGroupByResult(c, results) + if err != nil { + return nil, errors.Wrap(err, "applying limit/offset") + } + } + + // TODO as an optimization, we could apply some "having" + // conditions here long as they aren't on the Count(Distinct) + // aggregate + + // Calculate Count(Distinct) aggregate if requested. + aggregate, _, err := c.CallArg("aggregate") + if err == nil && aggregate != nil && aggregate.Name == "Count" && len(aggregate.Children) > 0 && aggregate.Children[0].Name == "Distinct" && !opt.Remote { + for n, gc := range results { + intersectRows := make([]*pql.Call, 0, len(gc.Group)) + for _, fr := range gc.Group { + var value interface{} = fr.RowID + // use fr.Value instead of fr.RowID if set (from int fields) + if fr.Value != nil { + value = &pql.Condition{Op: pql.EQ, Value: *fr.Value} + } + intersectRows = append(intersectRows, &pql.Call{Name: "Row", Args: map[string]interface{}{fr.Field: value}}) + } + // apply any filter, if present + if filter != nil { + intersectRows = append(intersectRows, filter) + } + // also intersect with any children of Distinct + if len(aggregate.Children[0].Children) > 0 { + intersectRows = append(intersectRows, aggregate.Children[0].Children[0]) + } + + countDistinctIntersect := &pql.Call{ + Name: "Count", + Children: []*pql.Call{ + { + Name: "Distinct", + Children: []*pql.Call{ + { + Name: "Intersect", + Children: intersectRows, + }, + }, + Args: aggregate.Children[0].Args, + Type: pql.PrecallGlobal, + }, + }, + } + + opt.PreTranslated = true + aggregateCount, err := o.execute(ctx, index, &pql.Query{Calls: []*pql.Call{countDistinctIntersect}}, []uint64{}, opt) + if err != nil { + return nil, err + } + results[n].Agg = int64(aggregateCount[0].(uint64)) + } + } + + // Apply having. + if hasHaving && !opt.Remote { + // parse the condition as PQL + if having.Name != "Condition" { + return nil, errors.New(errors.ErrUncoded, "the only supported having call is Condition()") + } + if len(having.Args) != 1 { + return nil, errors.New(errors.ErrUncoded, "Condition() must contain a single condition") + } + for subj, cond := range having.Args { + switch subj { + case "count", "sum": + results = featurebase.ApplyConditionToGroupCounts(results, subj, cond.(*pql.Condition)) + default: + return nil, errors.New(errors.ErrUncoded, "Condition() only supports count or sum") + } + } + } + + if sorter != nil && !opt.Remote { + sorter.data = results + sort.Stable(sorter) + results, err = applyLimitAndOffsetToGroupByResult(c, results) + if err != nil { + return nil, errors.Wrap(err, "applying limit/offset") + } + } else if hasHaving && !opt.Remote { + results, err = applyLimitAndOffsetToGroupByResult(c, results) + if err != nil { + return nil, errors.Wrap(err, "applying limit/offset") + } + + } + + aggType := "" + if aggregate != nil { + switch aggregate.Name { + case "Sum": + aggType = "sum" + case "Count": + aggType = "aggregate" + } + } + for _, res := range results { + if res.DecimalAgg != nil && aggType == "sum" { + aggType = "decimalSum" + break + } + } + + return featurebase.NewGroupCounts(aggType, results...), nil +} + +func applyLimitAndOffsetToGroupByResult(c *pql.Call, results []featurebase.GroupCount) ([]featurebase.GroupCount, error) { + // Apply offset. + if offset, hasOffset, err := c.UintArg("offset"); err != nil { + return nil, err + } else if hasOffset { + if int(offset) < len(results) { + results = results[offset:] + } + } + // Apply limit. + if limit, hasLimit, err := c.UintArg("limit"); err != nil { + return nil, err + } else if hasLimit { + if int(limit) < len(results) { + results = results[:limit] + } + } + return results, nil +} + +// mergeGroupCounts merges two slices of GroupCounts throwing away any that go +// beyond the limit. It assume that the two slices are sorted by the row ids in +// the fields of the group counts. It may modify its arguments. +func mergeGroupCounts(a, b []featurebase.GroupCount, limit int) []featurebase.GroupCount { + if limit > len(a)+len(b) { + limit = len(a) + len(b) + } + ret := make([]featurebase.GroupCount, 0, limit) + i, j := 0, 0 + for i < len(a) && j < len(b) && len(ret) < limit { + switch a[i].Compare(b[j]) { + case -1: + ret = append(ret, a[i]) + i++ + case 0: + a[i].Count += b[j].Count + a[i].Agg += b[j].Agg + if a[i].DecimalAgg != nil && b[j].DecimalAgg != nil { + sum := pql.AddDecimal(*a[i].DecimalAgg, *b[j].DecimalAgg) + a[i].DecimalAgg = &sum + } + ret = append(ret, a[i]) + i++ + j++ + case 1: + ret = append(ret, b[j]) + j++ + } + } + for ; i < len(a) && len(ret) < limit; i++ { + ret = append(ret, a[i]) + } + for ; j < len(b) && len(ret) < limit; j++ { + ret = append(ret, b[j]) + } + return ret +} + +func (o *orchestrator) executeRows(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (featurebase.RowIDs, error) { + // Fetch field name from argument. + // Check "field" first for backwards compatibility. + // TODO: remove at Pilosa 2.0 + var fieldName string + var ok bool + if fieldName, ok = c.Args["field"].(string); ok { + c.Args["_field"] = fieldName + } + if fieldName, ok = c.Args["_field"].(string); !ok { + return nil, errors.New(errors.ErrUncoded, "Rows() field required") + } + + // TODO(tlt): this is here to prevent the linter from complaining. + // Presumably this fieldName is/was used in code which is no longer here or + // is currently commented out. + _ = fieldName + + if columnID, ok, err := c.UintArg("column"); err != nil { + return nil, errors.Wrap(err, "getting column") + } else if ok { + shards = []uint64{columnID / featurebase.ShardWidth} + } + + // TODO, support "in" in conjunction w/ other args... or at least error if they're present together + if ids, found, err := c.UintSliceArg("in"); err != nil { + return nil, errors.Wrapf(err, "'in' argument of Rows must be a slice") + } else if found { + // "in" not supported with other args, so check here + for arg := range c.Args { + if arg != "field" && arg != "_field" && arg != "in" { + return nil, errors.Errorf("Rows call with 'in' does not support other arguments, but found '%s'", arg) + } + } + return ids, nil + } + + // Determine limit so we can use it when reducing. + limit := int(^uint(0) >> 1) + if lim, hasLimit, err := c.UintArg("limit"); err != nil { + return nil, err + } else if hasLimit { + limit = int(lim) + } + + // Merge returned results at coordinating node. + reduceFn := func(ctx context.Context, prev, v interface{}) interface{} { + other, _ := prev.(featurebase.RowIDs) + if err := ctx.Err(); err != nil { + return err + } + return other.Merge(v.(featurebase.RowIDs), limit) + } + // Get full result set. + other, err := o.mapReduce(ctx, index, shards, c, opt, reduceFn) + if err != nil { + return nil, err + } + results, _ := other.(featurebase.RowIDs) + + // TODO(jaffee) enable "like" support + // if !opt.Remote { + // if like, hasLike, err := c.StringArg("like"); err != nil { + // return nil, errors.Wrap(err, "getting like pattern") + // } else if hasLike { + // matches, err := e.Cluster.matchField(ctx, e.Holder.Field(index, fieldName), like) + // if err != nil { + // return nil, errors.Wrap(err, "matching like pattern") + // } + + // i, j, k := 0, 0, 0 + // for i < len(results) && j < len(matches) { + // x, y := results[i], matches[j] + // switch { + // case x < y: + // i++ + // case y < x: + // j++ + // default: + // results[k] = x + // i++ + // j++ + // k++ + // } + // } + // results = results[:k] + // } + // } + + return results, nil +} + +func (o *orchestrator) executeExtract(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (featurebase.ExtractedIDMatrix, error) { + // Extract the column filter call. + if len(c.Children) < 1 { + return featurebase.ExtractedIDMatrix{}, errors.New(errors.ErrUncoded, "missing column filter in Extract") + } + + // Extract fields from rows calls. + fields := make([]string, len(c.Children)-1) + for i, rows := range c.Children[1:] { + if rows.Name != "Rows" { + return featurebase.ExtractedIDMatrix{}, errors.Errorf("child call of Extract is %q but expected Rows", rows.Name) + } + var fieldName string + var ok bool + for k, v := range rows.Args { + switch k { + case "field", "_field": + fieldName = v.(string) + ok = true + default: + return featurebase.ExtractedIDMatrix{}, errors.Errorf("unsupported Rows argument for Extract: %q", k) + } + } + if !ok { + return featurebase.ExtractedIDMatrix{}, errors.New(errors.ErrUncoded, "missing field specification in Rows") + } + fields[i] = fieldName + } + + // Merge returned results at coordinating node. + reduceFn := func(ctx context.Context, prev, v interface{}) interface{} { + other, _ := prev.(featurebase.ExtractedIDMatrix) + if err := ctx.Err(); err != nil { + return err + } + other.Append(v.(featurebase.ExtractedIDMatrix)) + return other + } + + // Get full result set. + other, err := o.mapReduce(ctx, index, shards, c, opt, reduceFn) + if err != nil { + return featurebase.ExtractedIDMatrix{}, err + } + results, _ := other.(featurebase.ExtractedIDMatrix) + sort.Slice(results.Columns, func(i, j int) bool { + return results.Columns[i].ColumnID < results.Columns[j].ColumnID + }) + return results, nil +} + +func (o *orchestrator) executeConstRow(ctx context.Context, index string, c *pql.Call) (res *featurebase.Row, err error) { + // Fetch user-provided columns list. + ids, ok := c.Args["columns"].([]uint64) + if !ok { + return nil, errors.New(errors.ErrUncoded, "missing columns list") + } + + return featurebase.NewRow(ids...), nil +} + +func (o *orchestrator) executeUnionRows(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (*featurebase.Row, error) { + // Turn UnionRows(Rows(...)) into Union(Row(...), ...). + var rows []*pql.Call + for _, child := range c.Children { + // Check that we can use the call. + switch child.Name { + case "Rows": + case "TopN": + default: + return nil, errors.Errorf("cannot use %v as a rows query", child) + } + + // Execute the call. + rowsResult, err := o.executeCall(ctx, index, child, shards, opt) + if err != nil { + return nil, err + } + + // Turn the results into rows calls. + var resultRows []*pql.Call + switch rowsResult := rowsResult.(type) { + case *featurebase.PairsField: + // Translate pairs into rows calls. + for _, p := range rowsResult.Pairs { + var val interface{} + switch { + case p.Key != "": + val = p.Key + default: + val = p.ID + } + resultRows = append(resultRows, &pql.Call{ + Name: "Row", + Args: map[string]interface{}{ + rowsResult.Field: val, + }, + }) + } + case featurebase.RowIDs: + // Translate Row IDs into Row calls. + for _, id := range rowsResult { + resultRows = append(resultRows, &pql.Call{ + Name: "Row", + Args: map[string]interface{}{ + child.Args["_field"].(string): id, + }, + }) + } + default: + return nil, errors.Errorf("unexpected Rows type %T", rowsResult) + } + + // Propogate any special properties of the call. + switch child.Name { + case "Rows": + // Propogate "from" time, if set. + if v, ok := child.Args["from"]; ok { + for _, rowCall := range resultRows { + rowCall.Args["from"] = v + } + } + + // Propogate "to" time, if set. + if v, ok := child.Args["to"]; ok { + for _, rowCall := range resultRows { + rowCall.Args["to"] = v + } + } + } + + rows = append(rows, resultRows...) + } + + // Generate a Union call over the rows. + c = &pql.Call{ + Name: "Union", + Children: rows, + } + + // Execute the generated Union() call. + return o.executeBitmapCall(ctx, index, c, shards, opt) +} + +// executeCount executes a count() call. +func (o *orchestrator) executeCount(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *featurebase.ExecOptions) (uint64, error) { + span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeCount") + defer span.Finish() + + if len(c.Children) == 0 { + return 0, errors.New(errors.ErrUncoded, "Count() requires an input bitmap") + } else if len(c.Children) > 1 { + return 0, errors.New(errors.ErrUncoded, "Count() only accepts a single bitmap input") + } + + child := c.Children[0] + + // If the child is distinct/similar, execute it directly here and count the result. + if child.Type == pql.PrecallGlobal { + result, err := o.executeCall(ctx, index, child, shards, opt) + if err != nil { + return 0, err + } + + switch row := result.(type) { + case *featurebase.Row: + return row.Count(), nil + case featurebase.SignedRow: + return row.Pos.Count() + row.Neg.Count(), nil + case featurebase.DistinctTimestamp: + return uint64(len(row.Values)), nil + default: + return 0, errors.Errorf("cannot count result of type %T from call %q", row, child.String()) + } + } + + // Merge returned results at coordinating node. + reduceFn := func(ctx context.Context, prev, v interface{}) interface{} { + other, _ := prev.(uint64) + return other + v.(uint64) + } + + result, err := o.mapReduce(ctx, index, shards, c, opt, reduceFn) + if err != nil { + return 0, err + } + n, _ := result.(uint64) + + return n, nil +} + +// remoteExec executes a PQL query remotely for a set of shards on a node. +func (o *orchestrator) remoteExec(ctx context.Context, node dax.Address, index string, q *pql.Query, shards []uint64, embed []*featurebase.Row) (results []interface{}, err error) { // nolint: interfacer + span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeExec") + defer span.Finish() + + // Encode request object. + pbreq := &featurebase.QueryRequest{ + Query: q.String(), + Shards: shards, + Remote: true, + EmbeddedData: embed, + } + + scheme := node.Scheme() + if scheme == "" { + scheme = "http" + } + resp, err := o.client.QueryNode(ctx, &net.URI{ + Scheme: scheme, + Host: node.Host(), + Port: node.Port(), + }, index, pbreq) + if err != nil { + return nil, err + } + + return resp.Results, resp.Err +} + +// mapReduce maps and reduces data across the cluster. +// +// If a mapping of shards to a node fails then the shards are resplit across +// secondary nodes and retried. This continues to occur until all nodes are exhausted. +// +// mapReduce has to ensure that it never returns before any work it spawned has +// terminated. It's not enough to cancel the jobs; we have to wait for them to be +// done, or we can unmap resources they're still using. +func (o *orchestrator) mapReduce(ctx context.Context, index string, shards []uint64, c *pql.Call, opt *featurebase.ExecOptions, reduceFn reduceFunc) (result interface{}, err error) { + span, ctx := tracing.StartSpanFromContext(ctx, "Executor.mapReduce") + defer span.Finish() + + ch := make(chan mapResponse) + + // Wrap context with a cancel to kill goroutines on exit. + ctx, cancel := context.WithCancel(ctx) + // Create an errgroup so we can wait for all the goroutines to exit + eg, ctx := errgroup.WithContext(ctx) + + // After we're done processing, we have to wait for any outstanding + // functions in the ErrGroup to complete. If we didn't have an error + // already at that point, we'll report any errors from the ErrGroup + // instead. + defer func() { + cancel() + errWait := eg.Wait() + if err == nil { + err = errWait + } + }() + + nodes, err := o.topology.ComputeNodes(ctx, index, shards) + if err != nil { + return nil, errors.Wrapf(err, "getting nodes/shards for index '%q'", index) + } + + // Start mapping across all primary owners. + if err = o.mapper(ctx, eg, ch, index, nodes, c, opt, reduceFn); err != nil { + return nil, errors.Wrap(err, "starting mapper") + } + + // Iterate over all map responses and reduce. + expected := 0 + for _, n := range nodes { + expected += len(n.Shards) + } + done := ctx.Done() + for expected > 0 { + select { + case <-done: + return nil, ctx.Err() + case resp := <-ch: + if resp.err != nil { + cancel() // TODO(jaffee) I added this... seems right, but wasn't there before + return nil, errors.Wrap(resp.err, "mapping on primary node") + } + // if we got a response that we aren't discarding + // because it's an error, subtract it from our count... + expected -= len(resp.shards) + + // Reduce value. + result = reduceFn(ctx, result, resp.result) + var ok bool + // note *not* shadowed. + if err, ok = result.(error); ok { + cancel() + return nil, err + } + } + } + // note the deferred Wait above which might override this nil. + return result, nil +} + +// makeEmbeddedDataForShards produces new rows containing the RowSegments +// that would correspond to a given set of shards. +func makeEmbeddedDataForShards(allRows []*featurebase.Row, shards []uint64) []*featurebase.Row { + if len(allRows) == 0 || len(shards) == 0 { + return nil + } + newRows := make([]*featurebase.Row, len(allRows)) + for i, row := range allRows { + if row == nil || len(row.Segments) == 0 { + continue + } + if row.NoSplit { + newRows[i] = row + continue + } + segments := row.Segments + segmentIndex := 0 + newRows[i] = &featurebase.Row{ + Index: row.Index, + Field: row.Field, + } + for _, shard := range shards { + for segmentIndex < len(segments) && segments[segmentIndex].Shard() < shard { + segmentIndex++ + } + // no more segments in this row + if segmentIndex >= len(segments) { + break + } + if segments[segmentIndex].Shard() == shard { + newRows[i].Segments = append(newRows[i].Segments, segments[segmentIndex]) + segmentIndex++ + if segmentIndex >= len(segments) { + // no more segments, we're done + break + } + } + // if we got here, segments[segmentIndex].shard exists + // but is greater than the current shard, so we continue. + } + } + return newRows +} + +func (o *orchestrator) mapper(ctx context.Context, eg *errgroup.Group, ch chan mapResponse, index string, nodes []controller.ComputeNode, c *pql.Call, opt *featurebase.ExecOptions, reduceFn reduceFunc) (reterr error) { + span, ctx := tracing.StartSpanFromContext(ctx, "Executor.mapper") + defer span.Finish() + + // Group shards together by nodes. + done := ctx.Done() + + // Execute each node in a separate goroutine. + for _, node := range nodes { + node := node + shards := make([]uint64, len(node.Shards)) + for i, dshard := range node.Shards { + shards[i] = uint64(dshard) + } + eg.Go(func() error { + resp := mapResponse{node: node.Address, shards: shards} + + var embeddedRowsForNode []*featurebase.Row + if opt.EmbeddedData != nil { + embeddedRowsForNode = makeEmbeddedDataForShards(opt.EmbeddedData, shards) + } + + attempts := 0 + for ; attempts == 0 || (resp.err != nil && strings.Contains(resp.err.Error(), errConnectionRefused) && attempts < 3); attempts++ { + // On error retry against remaining nodes. If an error returns then + // the context will cancel and cause all open goroutines to return. + // + // We distinguish here between an error which indicates that the + // node is not available (and therefore we need to failover to a + // replica) and a valid error from a healthy node. In the case of + // the latter, there's no need to retry a replica, we should trust + // the error from the healthy node and return that immediately. + // TODO(jaffee) retries should contact MDS and find out who is up and has access to shards needed + results, err := o.remoteExec(ctx, node.Address, index, &pql.Query{Calls: []*pql.Call{c}}, shards, embeddedRowsForNode) + if len(results) > 0 { + resp.result = results[0] + } + resp.err = err + } + // Return response to the channel. + select { + case <-done: + // If someone just canceled the context + // arbitrarily, we could end up here with this + // being the first non-nil error handed to + // the ErrGroup, in which case, it's the best + // explanation we have for why everything's + // stopping. + return ctx.Err() + case ch <- resp: + return nil + } + }) + if reterr != nil { + return reterr // exit early if error occurs when running serially + } + } + return nil +} + +func (o *orchestrator) preTranslate(ctx context.Context, index string, calls ...*pql.Call) (cols map[string]map[string]uint64, rows map[string]map[string]map[string]uint64, err error) { + // Collect all of the required keys. + collector := keyCollector{ + createCols: make(map[string][]string), + findCols: make(map[string][]string), + createRows: make(map[string]map[string][]string), + findRows: make(map[string]map[string][]string), + } + for _, call := range calls { + err := o.collectCallKeys(&collector, call, index) + if err != nil { + return nil, nil, err + } + } + + // Create keys. + // Both rows and columns need to be created first because of foreign index keys. + cols = make(map[string]map[string]uint64) + rows = make(map[string]map[string]map[string]uint64) + for index, keys := range collector.createCols { + translations, err := o.trans.CreateIndexKeys(ctx, index, keys) + if err != nil { + return nil, nil, errors.Wrap(err, "creating query column keys") + } + cols[index] = translations + } + for index, fields := range collector.createRows { + idxRows := make(map[string]map[string]uint64) + for field, keys := range fields { + translations, err := o.trans.CreateFieldKeys(ctx, index, field, keys) + if err != nil { + return nil, nil, errors.Wrap(err, "creating query row keys") + } + idxRows[field] = translations + } + rows[index] = idxRows + } + + // Find other keys. + for index, keys := range collector.findCols { + translations, err := o.trans.FindIndexKeys(ctx, index, keys) + if err != nil { + return nil, nil, errors.Wrap(err, "finding query column keys") + } + if prev := cols[index]; prev != nil { + for key, id := range translations { + prev[key] = id + } + } else { + cols[index] = translations + } + } + for index, fields := range collector.findRows { + idxRows := rows[index] + if idxRows == nil { + idxRows = make(map[string]map[string]uint64) + rows[index] = idxRows + } + for field, keys := range fields { + translations, err := o.trans.FindFieldKeys(ctx, index, field, keys) + if err != nil { + return nil, nil, errors.Wrap(err, "finding query row keys") + } + if prev := idxRows[field]; prev != nil { + for key, id := range translations { + prev[key] = id + } + } else { + idxRows[field] = translations + } + } + } + + return cols, rows, nil +} + +func (o *orchestrator) collectCallKeys(dst *keyCollector, c *pql.Call, index string) error { + // Check for an overriding 'index' argument. + // This also applies to all child calls. + if callIndex := c.CallIndex(); callIndex != "" { + index = callIndex + } + + // Handle the field arg. + switch c.Name { + case "Set": + if field, err := c.FieldArg(); err == nil { + if arg, ok := c.Args[field].(string); ok { + dst.CreateRows(index, field, arg) + } + } + + // TODO: will have to consider how to support Store... creating the field if it doesn't exist will not be a thing though. + case "Store": + return errors.New(errors.ErrUncoded, "Store query currently unsupported") + case "Clear", "Row", "Range", "ClearRow": + if field, err := c.FieldArg(); err == nil { + switch arg := c.Args[field].(type) { + case string: + dst.FindRows(index, field, arg) + case *pql.Condition: + // This is a workaround to allow `==` and `!=` to work on foreign index fields. + if key, ok := arg.Value.(string); ok { + switch arg.Op { + case pql.EQ, pql.NEQ: + dst.FindRows(index, field, key) + default: + return errors.Errorf("operator %v not defined on strings", arg.Op) + } + } + } + } + } + + // Handle _col. + if col, ok := c.Args["_col"].(string); ok { + switch c.Name { + case "Set": + dst.CreateColumns(index, col) + default: + dst.FindColumns(index, col) + } + } + + // Handle _row. + if row, ok := c.Args["_row"].(string); ok { + // Find the field. + field, ok, err := c.StringArg("_field") + if err != nil { + return errors.Wrap(err, "finding field") + } + if !ok { + return errors.Wrap(ErrFieldNotFound, "finding field for _row argument") + } + + dst.FindRows(index, field, row) + } + + // Handle queries that need a "column" argument. + switch c.Name { + case "Rows", "GroupBy", "FieldValue", "IncludesColumn": + if col, ok := c.Args["column"].(string); ok { + dst.FindColumns(index, col) + } + } + + // Handle special per-query arguments. + switch c.Name { + case "ConstRow": + // Translate the columns list. + if cols, ok := c.Args["columns"].([]interface{}); ok { + keys := make([]string, 0, len(cols)) + for _, v := range cols { + switch v := v.(type) { + case string: + keys = append(keys, v) + case uint64: + case int64: + default: + return errors.Errorf("invalid column identifier %v of type %T", c, c) + } + } + dst.FindColumns(index, keys...) + } + + case "Rows": + // Find the field. + var field string + if f, ok1, err := c.StringArg("_field"); err != nil { + return errors.Wrap(err, "finding _field for Rows previous translation") + } else if ok1 { + field = f + } else if f, ok2, err := c.StringArg("field"); err != nil { + return errors.Wrap(err, "finding field for Rows previous translation") + } else if ok2 { + field = f + } else { + return errors.New(errors.ErrUncoded, "missing field in Rows call") + } + if prev, ok := c.Args["previous"].(string); ok { + dst.FindRows(index, field, prev) + } + if in, ok := c.Args["in"]; ok { + inIn, ok := in.([]interface{}) + if !ok { + return errors.Errorf("unexpected type for argument 'in' %v of %[1]T", inIn) + } + inStrs := make([]string, 0) + for _, v := range inIn { + if vstr, ok := v.(string); ok { + inStrs = append(inStrs, vstr) + } + } + dst.FindRows(index, field, inStrs...) + } + } + + // Collect keys from child calls. + for _, child := range c.Children { + err := o.collectCallKeys(dst, child, index) + if err != nil { + return err + } + } + + // Collect keys from argument calls. + for _, arg := range c.Args { + argCall, ok := arg.(*pql.Call) + if !ok { + continue + } + + err := o.collectCallKeys(dst, argCall, index) + if err != nil { + return err + } + } + + return nil +} + +type keyCollector struct { + createCols, findCols map[string][]string // map[index] -> column keys + createRows, findRows map[string]map[string][]string // map[index]map[field] -> row keys +} + +func (c *keyCollector) CreateColumns(index string, columns ...string) { + if len(columns) == 0 { + return + } + c.createCols[index] = append(c.createCols[index], columns...) +} + +func (c *keyCollector) FindColumns(index string, columns ...string) { + if len(columns) == 0 { + return + } + c.findCols[index] = append(c.findCols[index], columns...) +} + +func (c *keyCollector) CreateRows(index string, field string, columns ...string) { + if len(columns) == 0 { + return + } + idx := c.createRows[index] + if idx == nil { + idx = make(map[string][]string) + c.createRows[index] = idx + } + idx[field] = append(idx[field], columns...) +} + +func (c *keyCollector) FindRows(index string, field string, columns ...string) { + if len(columns) == 0 { + return + } + idx := c.findRows[index] + if idx == nil { + idx = make(map[string][]string) + c.findRows[index] = idx + } + idx[field] = append(idx[field], columns...) +} + +func fieldValidateValue(f *featurebase.FieldInfo, val interface{}) error { + if val == nil { + return nil + } + + // Validate special types. + switch val := val.(type) { + case string: + if !f.Options.Keys { + return errors.Errorf("string value on unkeyed field %q", f.Name) + } + return nil + case *pql.Condition: + switch v := val.Value.(type) { + case nil: + case string: + case uint64: + case int64: + case float64: + case pql.Decimal: + case time.Time: + case []interface{}: + for _, v := range v { + if err := fieldValidateValue(f, v); err != nil { + return err + } + } + return nil + default: + return errors.Errorf("invalid value %v in condition %q", v, val.String()) + } + return fieldValidateValue(f, val.Value) + } + + switch f.Options.Type { + case FieldTypeSet, FieldTypeMutex, FieldTypeTime: + switch v := val.(type) { + case uint64: + case int64: + if v < 0 { + return errors.Errorf("negative ID %d for set field %q", v, f.Name) + } + default: + return errors.Errorf("invalid value %v for field %q of type %s", v, f.Name, f.Options.Type) + } + if f.Options.Keys { + return errors.Errorf("found integer ID %d on keyed field %q", val, f.Name) + } + case FieldTypeBool: + switch v := val.(type) { + case bool: + default: + return errors.Errorf("invalid value %v for bool field %q", v, f.Name) + } + case FieldTypeInt: + switch v := val.(type) { + case uint64: + if v > 1<<63 { + return errors.Errorf("oversized integer %d for int field %q (range: -2^63 to 2^63-1)", v, f.Name) + } + case int64: + default: + return errors.Errorf("invalid value %v for int field %q", v, f.Name) + } + case FieldTypeDecimal: + switch v := val.(type) { + case uint64: + case int64: + case float64: + case pql.Decimal: + default: + return errors.Errorf("invalid value %v for decimal field %q", v, f.Name) + } + case FieldTypeTimestamp: + switch v := val.(type) { + case time.Time: + default: + return errors.Errorf("invalid value %v for timestamp field %q", v, f.Name) + } + default: + return errors.Errorf("unsupported type %s of field %q", f.Options.Type, f.Name) + } + + return nil +} + +func (o *orchestrator) translateCall(ctx context.Context, c *pql.Call, index string, columnKeys map[string]map[string]uint64, rowKeys map[string]map[string]map[string]uint64) (*pql.Call, error) { + // Check for an overriding 'index' argument. + // This also applies to all child calls. + if callIndex := c.CallIndex(); callIndex != "" { + index = callIndex + } + idx, err := o.schema.IndexInfo(ctx, index) + if err != nil { + return nil, errors.Wrapf(err, "translating query on index %q", index) + } + + // Fetch the column keys list for this index. + indexCols, indexRows := columnKeys[index], rowKeys[index] + + // Handle the field arg. + switch c.Name { + case "Set", "Store": + if field, err := c.FieldArg(); err == nil { + f, err := o.schema.FieldInfo(ctx, index, field) + if err != nil { + return nil, errors.Wrapf(err, "validating value for field %q", field) + } + arg := c.Args[field] + if err := fieldValidateValue(f, arg); err != nil { + return nil, errors.Wrap(err, "validating store value") + } + switch arg := arg.(type) { + case string: + if translation, ok := indexRows[field][arg]; ok { + c.Args[field] = translation + } else { + return nil, errors.Wrapf(featurebase.ErrTranslatingKeyNotFound, "destination key not found %q in %q in index %q", arg, field, index) + } + case bool: + if arg { + c.Args[field] = trueRowID + } else { + c.Args[field] = falseRowID + } + } + } + + case "Clear", "Row", "Range", "ClearRow": + if field, err := c.FieldArg(); err == nil { + f, err := o.schema.FieldInfo(ctx, index, field) + if err != nil { + return nil, errors.Wrapf(err, "validating value for field %q", field) + } + arg := c.Args[field] + if err := fieldValidateValue(f, arg); err != nil { + return nil, errors.Wrap(err, "validating field parameter value") + } + if c.Name == "Row" { + switch f.Options.Type { + case FieldTypeInt, FieldTypeDecimal, FieldTypeTimestamp: + if _, ok := arg.(*pql.Condition); !ok { + // This is workaround to support pql.ASSIGN ('=') as condition ('==') for BSI fields. + arg = &pql.Condition{ + Op: pql.EQ, + Value: arg, + } + c.Args[field] = arg + } + } + } + switch arg := arg.(type) { + case string: + if translation, ok := indexRows[field][arg]; ok { + c.Args[field] = translation + } else { + // Rewrite the call into a zero value call. + return o.callZero(c), nil + } + case bool: + if arg { + c.Args[field] = trueRowID + } else { + c.Args[field] = falseRowID + } + case *pql.Condition: + // This is a workaround to allow `==` and `!=` to work on foreign index fields. + if key, ok := arg.Value.(string); ok { + switch arg.Op { + case pql.EQ, pql.NEQ: + if translation, ok := indexRows[field][key]; ok { + arg.Value = translation + } else { + // Rewrite the call into a zero value call. + return o.callZero(c), nil + } + default: + return nil, errors.Errorf("operator %v not defined on strings", arg.Op) + } + } + } + } + } + + // Handle _col. + if col, ok := c.Args["_col"].(string); ok { + if !idx.Options.Keys { + return nil, errors.Wrapf(featurebase.ErrTranslatingKeyNotFound, "translating column on unkeyed index %q", index) + } + if id, ok := indexCols[col]; ok { + c.Args["_col"] = id + } else { + switch c.Name { + case "Set": + return nil, errors.Wrapf(featurebase.ErrTranslatingKeyNotFound, "destination key not found %q in index %q", col, index) + default: + return o.callZero(c), nil + } + } + } + + // Handle _row. + if row, ok := c.Args["_row"]; ok { + // Find the field. + var field string + if f, ok1, err := c.StringArg("_field"); err != nil { + return nil, errors.Wrap(err, "finding _field") + } else if ok1 { + field = f + } else if f, ok2, err := c.StringArg("field"); err != nil { + return nil, errors.Wrap(err, "finding field") + } else if ok2 { + field = f + } else { + return nil, errors.New(errors.ErrUncoded, "missing field") + } + + f, err := o.schema.FieldInfo(ctx, index, field) + if err != nil { + return nil, errors.Wrapf(err, "validating value for field %q", field) + } + if err := fieldValidateValue(f, row); err != nil { + return nil, errors.Wrap(err, "validating row value") + } + switch row := row.(type) { + case string: + if translation, ok := indexRows[field][row]; ok { + c.Args["_row"] = translation + } else { + return o.callZero(c), nil + } + } + } + + // Handle queries that need a "column" argument. + switch c.Name { + case "Rows", "GroupBy", "FieldValue", "IncludesColumn": + if col, ok := c.Args["column"].(string); ok { + if translation, ok := indexCols[col]; ok { + c.Args["column"] = translation + } else { + // Rewrite the call into a zero value call. + return o.callZero(c), nil + } + } + } + + // Handle special per-query arguments. + switch c.Name { + case "ConstRow": + // Translate the columns list. + if cols, ok := c.Args["columns"].([]interface{}); ok { + out := make([]uint64, 0, len(cols)) + for _, v := range cols { + switch v := v.(type) { + case string: + if id, ok := indexCols[v]; ok { + out = append(out, id) + } + case uint64: + out = append(out, v) + case int64: + out = append(out, uint64(v)) + default: + return nil, errors.Errorf("invalid column identifier %v of type %T", c, c) + } + } + c.Args["columns"] = out + } + + case "Rows": + // Find the field. + var field string + if f, ok1, err := c.StringArg("_field"); err != nil { + return nil, errors.Wrap(err, "finding _field for Rows previous translation") + } else if ok1 { + field = f + } else if f, ok2, err := c.StringArg("field"); err != nil { + return nil, errors.Wrap(err, "finding field for Rows previous translation") + } else if ok2 { + field = f + } else { + return nil, errors.New(errors.ErrUncoded, "missing field in Rows call") + } + // Translate the previous row key. + if prev, ok := c.Args["previous"]; ok { + // Validate the type. + f, err := o.schema.FieldInfo(ctx, index, field) + if err != nil { + return nil, errors.Wrapf(err, "validating value for field %q", field) + } + if err := fieldValidateValue(f, prev); err != nil { + return nil, errors.Wrap(err, "validating prev value") + } + + switch prev := prev.(type) { + case string: + // Look up a translation for the previous row key. + if translation, ok := indexRows[field][prev]; ok { + c.Args["previous"] = translation + } else { + return nil, errors.Wrapf(featurebase.ErrTranslatingKeyNotFound, "translating previous key %q from field %q in index %q in Rows call", prev, field, index) + } + case bool: + if prev { + c.Args["previous"] = trueRowID + } else { + c.Args["previous"] = falseRowID + } + } + } + + // Check if "like" argument is applied to keyed fields. + if _, found := c.Args["like"].(string); found { + fieldName, err := c.FirstStringArg("_field", "field") + if err != nil || fieldName == "" { + return nil, fmt.Errorf("cannot read field name for Rows call") + } + if f, err := o.schema.FieldInfo(ctx, index, fieldName); err != nil { + return nil, errors.Wrapf(err, "getting field %q", fieldName) + } else if !f.Options.Keys { + return nil, fmt.Errorf("'%s' is not a set/mutex/time field with a string key", fieldName) + } + } + + if in, ok := c.Args["in"]; ok { + inIn, ok := in.([]interface{}) + if !ok { + return nil, errors.Errorf("unexpected type for argument 'in' %v of %[1]T", in) + } + inIDs := make([]interface{}, 0, len(inIn)) + for _, inVal := range inIn { + if inStr, ok := inVal.(string); ok { + id, found := rowKeys[index][field][inStr] + if found { + inIDs = append(inIDs, id) + } + } else { + inIDs = append(inIDs, inVal) + } + } + c.Args["in"] = inIDs + } + } + + // Translate child calls. + for i, child := range c.Children { + translated, err := o.translateCall(ctx, child, index, columnKeys, rowKeys) + if err != nil { + return nil, err + } + c.Children[i] = translated + } + + // Translate argument calls. + for k, arg := range c.Args { + argCall, ok := arg.(*pql.Call) + if !ok { + continue + } + + translated, err := o.translateCall(ctx, argCall, index, columnKeys, rowKeys) + if err != nil { + return nil, err + } + + c.Args[k] = translated + } + + return c, nil +} + +func (o *orchestrator) callZero(c *pql.Call) *pql.Call { + switch c.Name { + case "Row", "Range": + if field, err := c.FieldArg(); err == nil { + if cond, ok := c.Args[field].(*pql.Condition); ok { + if cond.Op == pql.NEQ { + // Turn not nothing into everything. + return &pql.Call{Name: "All"} + } + } + } + + // Use an empty union as a placeholder. + return &pql.Call{Name: "Union"} + + default: + return nil + } +} + +func (o *orchestrator) translateResults(ctx context.Context, index string, idx *featurebase.IndexInfo, calls []*pql.Call, results []interface{}, memoryAvailable int64) (err error) { + span, _ := tracing.StartSpanFromContext(ctx, "Executor.translateResults") + defer span.Finish() + + idMap := make(map[uint64]string) + if idx.Options.Keys { + // Collect all index ids. + idSet := make(map[uint64]struct{}) + for i := range calls { + if err := o.collectResultIDs(ctx, idx, calls[i], results[i], idSet); err != nil { + return err + } + } + if idMap, err = o.trans.TranslateIndexIDSet(ctx, index, idSet); err != nil { + return err + } + } + + for i := range results { + results[i], err = o.translateResult(ctx, idx, calls[i], results[i], idMap) + if err != nil { + return err + } + } + return nil +} + +// translationStrategy denotes the several different ways the bits in +// a *Row could be translated to string keys. +type translationStrategy int + +const ( + // byCurrentIndex means to interpret the bits as IDs in "top + // level" index for this query (e.g. the index specified in the + // path of the HTTP request). + byCurrentIndex translationStrategy = iota + 1 + // byRowField means that the bits in this *Row are row IDs which + // should be translated using the field's (*Row.Field) translation store. + byRowField + // byRowFieldForeignIndex means that the bits in this *Row should + // be interpreted as IDs in the foreign index of the *Row.Field. + byRowFieldForeignIndex + // byRowIndex means the bits in this *Row should be translated + // according to the index named by *Row.Index + byRowIndex + // noTranslation means the bits should not be translated to string + // keys. + noTranslation +) + +// howToTranslate determines how a *Row object's bits should be +// translated to keys (if at all). There are several different options +// detailed by the various const values of translationStrategy. In +// order to do this it has to figure out the row's index and field +// which it also returns as the caller may need them to actually +// execute the translation or do whatever else it's doing with the +// translationStrategy information. +func (o *orchestrator) howToTranslate(ctx context.Context, idx *featurebase.IndexInfo, row *featurebase.Row) (rowIdx *featurebase.IndexInfo, rowField *featurebase.FieldInfo, strat translationStrategy, err error) { + // First get the index and field the row specifies (if any). + rowIdx = idx + if row.Index != "" && row.Index != idx.Name { + rowIdx, err = o.schema.IndexInfo(ctx, row.Index) + if err != nil { + return nil, nil, 0, errors.Wrapf(err, "got a row with unknown index: %s", row.Index) + } + } + if row.Field != "" { + rowField, err = o.schema.FieldInfo(ctx, row.Index, row.Field) + if err != nil { + return nil, nil, 0, errors.Wrapf(err, "got a row with unknown index/field %s/%s", idx.Name, row.Field) + } + } + + // Handle the case where the Row has specified a field. + if rowField != nil { + // Handle the case where field has a foreign index. + if rowField.Options.ForeignIndex != "" { + fidx, err := o.schema.IndexInfo(ctx, rowField.Options.ForeignIndex) + if err != nil { + return nil, nil, 0, errors.Errorf("foreign index %s not found for field %s in index %s", rowField.Options.ForeignIndex, rowField.Name, rowIdx.Name) + } + if fidx.Options.Keys { + return rowIdx, rowField, byRowFieldForeignIndex, nil + } + } else if rowField.Options.Keys { + return rowIdx, rowField, byRowField, nil + } + return rowIdx, rowField, noTranslation, nil + } + + // In this case, the row has specified an index, but not a field, + // so we translate according to that index. + if rowIdx != idx && rowIdx.Options.Keys { + return rowIdx, rowField, byRowIndex, nil + } + + // Handle the normal case (row represents a set of records in + // the top level index, Row has not specifed a different index + // or field). + if rowIdx == idx && idx.Options.Keys && rowField == nil { + return rowIdx, rowField, byCurrentIndex, nil + } + return rowIdx, rowField, noTranslation, nil +} + +func (o *orchestrator) collectResultIDs(ctx context.Context, idx *featurebase.IndexInfo, call *pql.Call, result interface{}, idSet map[uint64]struct{}) error { + switch result := result.(type) { + case *featurebase.Row: + // Only collect result IDs if they are in the current index. + _, _, strategy, err := o.howToTranslate(ctx, idx, result) + if err != nil { + return errors.Wrap(err, "determining how to translate") + } + if strategy == byCurrentIndex { + for _, segment := range result.Segments { + for _, col := range segment.Columns() { + idSet[col] = struct{}{} + } + } + } + case featurebase.ExtractedIDMatrix: + for _, col := range result.Columns { + idSet[col.ColumnID] = struct{}{} + } + } + + return nil +} + +// preTranslateMatrixSet translates the IDs of a set field in an extracted matrix. +func (o *orchestrator) preTranslateMatrixSet(ctx context.Context, mat featurebase.ExtractedIDMatrix, fieldIdx uint, index, field string) (map[uint64]string, error) { + ids := make(map[uint64]struct{}, len(mat.Columns)) + for _, col := range mat.Columns { + for _, v := range col.Rows[fieldIdx] { + ids[v] = struct{}{} + } + } + + return o.trans.TranslateFieldIDs(ctx, index, field, ids) +} + +func (o *orchestrator) translateResult(ctx context.Context, idx *featurebase.IndexInfo, call *pql.Call, result interface{}, idSet map[uint64]string) (_ interface{}, err error) { + switch result := result.(type) { + case *featurebase.Row: + rowIdx, rowField, strategy, err := o.howToTranslate(ctx, idx, result) + if err != nil { + return nil, errors.Wrap(err, "determining translation strategy") + } + switch strategy { + case byCurrentIndex: + other := &featurebase.Row{} + for _, segment := range result.Segments { + for _, col := range segment.Columns() { + other.Keys = append(other.Keys, idSet[col]) + } + } + return other, nil + case byRowField: + keys, err := o.trans.TranslateFieldListIDs(ctx, rowIdx.Name, rowField.Name, result.Columns()) + if err != nil { + return nil, errors.Wrap(err, "translating Row to field keys") + } + result.Keys = keys + case byRowFieldForeignIndex: + idx, err = o.schema.IndexInfo(ctx, rowField.Options.ForeignIndex) + if err != nil { + return nil, errors.Wrapf(err, "foreign index %s not found for field %s in index %s", rowField.Options.ForeignIndex, rowField.Name, rowIdx.Name) + } + for _, segment := range result.Segments { + keys, err := o.trans.TranslateIndexIDs(ctx, rowField.Options.ForeignIndex, segment.Columns()) + if err != nil { + return nil, errors.Wrap(err, "translating index ids") + } + result.Keys = append(result.Keys, keys...) + } + + case byRowIndex: + for _, segment := range result.Segments { + keys, err := o.trans.TranslateIndexIDs(ctx, rowIdx.Name, segment.Columns()) + if err != nil { + return nil, errors.Wrap(err, "translating index ids") + } + result.Keys = append(result.Keys, keys...) + } + return result, nil + + case noTranslation: + return result, nil + default: + return nil, errors.Errorf("unknown translation strategy %d", strategy) + } + case featurebase.SignedRow: + sr, err := func() (*featurebase.SignedRow, error) { + fieldName := callArgString(call, "field") + if fieldName == "" { + return nil, nil + } + + field, err := o.schema.FieldInfo(ctx, idx.Name, fieldName) + if err != nil { + return nil, nil + } + + if field.Options.Keys { + rslt := result.Pos + if rslt == nil { + return &featurebase.SignedRow{Pos: &featurebase.Row{}}, nil + } + other := &featurebase.Row{} + for _, segment := range rslt.Segments { + keys, err := o.trans.TranslateIndexIDs(ctx, field.Options.ForeignIndex, segment.Columns()) + if err != nil { + return nil, errors.Wrap(err, "translating index ids") + } + other.Keys = append(other.Keys, keys...) + } + return &featurebase.SignedRow{Pos: other}, nil + } + + return nil, nil + }() + if err != nil { + return nil, err + } else if sr != nil { + return *sr, nil + } + + case featurebase.PairField: + if fieldName := callArgString(call, "field"); fieldName != "" { + field, err := o.schema.FieldInfo(ctx, idx.Name, fieldName) + if err != nil { + return nil, fmt.Errorf("field %q not found", fieldName) + } + if field.Options.Keys { + // TODO(jaffee) get index name from call? CallIndex? (not just here) + keys, err := o.trans.TranslateFieldListIDs(ctx, idx.Name, fieldName, []uint64{result.Pair.ID}) + if err != nil { + return nil, err + } + key := keys[0] + if call.Name == "MinRow" || call.Name == "MaxRow" { + result.Pair.Key = key + return result, nil + } + return featurebase.PairField{ + Pair: featurebase.Pair{Key: key, Count: result.Pair.Count}, + Field: fieldName, + }, nil + } + } + + case *featurebase.PairsField: + if fieldName := callArgString(call, "_field"); fieldName != "" { + field, err := o.schema.FieldInfo(ctx, idx.Name, fieldName) + if err != nil { + return nil, errors.Wrapf(err, "field '%q'", fieldName) + } + if field.Options.Keys { + ids := make([]uint64, len(result.Pairs)) + for i := range result.Pairs { + ids[i] = result.Pairs[i].ID + } + keys, err := o.trans.TranslateFieldListIDs(ctx, idx.Name, fieldName, ids) + if err != nil { + return nil, err + } + other := make([]featurebase.Pair, len(result.Pairs)) + for i := range result.Pairs { + other[i] = featurebase.Pair{Key: keys[i], Count: result.Pairs[i].Count} + } + return &featurebase.PairsField{ + Pairs: other, + Field: fieldName, + }, nil + } + } + + case *featurebase.GroupCounts: + fieldIDs := make(map[*featurebase.FieldInfo]map[uint64]struct{}) + foreignIDs := make(map[*featurebase.FieldInfo]map[uint64]struct{}) + groups := result.Groups() + for _, gl := range groups { + for _, g := range gl.Group { + field, err := o.schema.FieldInfo(ctx, idx.Name, g.Field) + if err != nil { + return nil, errors.Wrapf(err, "getting field '%q", g.Field) + } + if field.Options.Keys { + if g.Value != nil { + if fi := field.Options.ForeignIndex; fi != "" { + m, ok := foreignIDs[field] + if !ok { + m = make(map[uint64]struct{}, len(groups)) + foreignIDs[field] = m + } + + m[uint64(*g.Value)] = struct{}{} + continue + } + } + + m, ok := fieldIDs[field] + if !ok { + m = make(map[uint64]struct{}, len(groups)) + fieldIDs[field] = m + } + + m[g.RowID] = struct{}{} + } + } + } + + fieldTranslations := make(map[string]map[uint64]string) + for field, ids := range fieldIDs { + trans, err := o.trans.TranslateFieldIDs(ctx, idx.Name, field.Name, ids) + if err != nil { + return nil, errors.Wrapf(err, "translating IDs in field '%q'", field.Name) + } + fieldTranslations[field.Name] = trans + } + + foreignTranslations := make(map[string]map[uint64]string) + for field, ids := range foreignIDs { + trans, err := o.trans.TranslateIndexIDSet(ctx, field.Options.ForeignIndex, ids) + if err != nil { + return nil, errors.Wrapf(err, "translating foreign IDs from index %q", field.Options.ForeignIndex) + } + foreignTranslations[field.Name] = trans + } + + // We are reluctant to smash result, and I'm not sure we need + // to be but I'm not sure we don't need to be. + newGroups := make([]featurebase.GroupCount, len(groups)) + copy(newGroups, groups) + for gi, gl := range groups { + + group := make([]featurebase.FieldRow, len(gl.Group)) + for i, g := range gl.Group { + if ft, ok := fieldTranslations[g.Field]; ok { + g.RowKey = ft[g.RowID] + } else if ft, ok := foreignTranslations[g.Field]; ok && g.Value != nil { + g.RowKey = ft[uint64(*g.Value)] + g.Value = nil + } + + group[i] = g + } + // Replace with translated group. + newGroups[gi].Group = group + } + if result != nil { + return featurebase.NewGroupCounts(result.AggregateColumn(), newGroups...), nil + } + return &featurebase.GroupCounts{}, nil + case featurebase.RowIDs: + fieldName := callArgString(call, "_field") + if fieldName == "" { + return nil, ErrFieldNotFound + } + + other := featurebase.RowIdentifiers{ + Field: fieldName, + } + + if field, err := o.schema.FieldInfo(ctx, idx.Name, fieldName); err != nil { + return nil, errors.Wrapf(err, "'%q'", fieldName) + } else if field.Options.Keys { + keys, err := o.trans.TranslateFieldListIDs(ctx, idx.Name, field.Name, result) + if err != nil { + return nil, errors.Wrap(err, "translating row IDs") + } + other.Keys = keys + } else { + other.Rows = result + } + + return other, nil + + case featurebase.ExtractedIDMatrix: + type fieldMapper = func([]uint64) (_ interface{}, err error) + + fields := make([]featurebase.ExtractedTableField, len(result.Fields)) + mappers := make([]fieldMapper, len(result.Fields)) + for i, v := range result.Fields { + field, err := o.schema.FieldInfo(ctx, idx.Name, v) + if err != nil { + return nil, errors.Wrapf(err, "'%q'", v) + } + + var mapper fieldMapper + var datatype string + switch typ := field.Options.Type; typ { + case FieldTypeBool: + datatype = "bool" + mapper = func(ids []uint64) (_ interface{}, err error) { + switch len(ids) { + case 0: + return nil, nil + case 1: + switch ids[0] { + case 0: + return false, nil + case 1: + return true, nil + default: + return nil, errors.Errorf("invalid ID for boolean %q: %d", field.Name, ids[0]) + } + default: + return nil, errors.Errorf("boolean %q has too many values: %v", field.Name, ids) + } + } + case FieldTypeSet, FieldTypeTime: + if field.Options.Keys { + datatype = "[]string" + translations, err := o.preTranslateMatrixSet(ctx, result, uint(i), idx.Name, field.Name) + if err != nil { + return nil, errors.Wrapf(err, "translating IDs of field %q", v) + } + mapper = func(ids []uint64) (interface{}, error) { + keys := make([]string, len(ids)) + for i, id := range ids { + keys[i] = translations[id] + } + return keys, nil + } + } else { + datatype = "[]uint64" + mapper = func(ids []uint64) (interface{}, error) { + if ids == nil { + ids = []uint64{} + } + return ids, nil + } + } + case FieldTypeMutex: + if field.Options.Keys { + datatype = "string" + translations, err := o.preTranslateMatrixSet(ctx, result, uint(i), idx.Name, field.Name) + if err != nil { + return nil, errors.Wrapf(err, "translating IDs of field %q", v) + } + mapper = func(ids []uint64) (interface{}, error) { + switch len(ids) { + case 0: + return nil, nil + case 1: + return translations[ids[0]], nil + default: + return nil, errors.Errorf("mutex %q has too many values: %v", field.Name, ids) + } + } + } else { + datatype = "uint64" + mapper = func(ids []uint64) (_ interface{}, err error) { + switch len(ids) { + case 0: + return nil, nil + case 1: + return ids[0], nil + default: + return nil, errors.Errorf("mutex %q has too many values: %v", field.Name, ids) + } + } + } + case FieldTypeInt: + if fi := field.Options.ForeignIndex; fi != "" { + if field.Options.Keys { + datatype = "string" + ids := make(map[uint64]struct{}, len(result.Columns)) + for _, col := range result.Columns { + for _, v := range col.Rows[i] { + ids[v] = struct{}{} + } + } + trans, err := o.trans.TranslateIndexIDSet(ctx, field.Options.ForeignIndex, ids) + if err != nil { + return nil, errors.Wrapf(err, "translating foreign IDs from index %q", field.Options.ForeignIndex) + } + mapper = func(ids []uint64) (interface{}, error) { + switch len(ids) { + case 0: + return nil, nil + case 1: + return trans[ids[0]], nil + default: + return nil, errors.Errorf("BSI field %q has too many values: %v", field.Name, ids) + } + } + } else { + datatype = "uint64" + mapper = func(ids []uint64) (interface{}, error) { + switch len(ids) { + case 0: + return nil, nil + case 1: + return ids[0], nil + default: + return nil, errors.Errorf("BSI field %q has too many values: %v", field.Name, ids) + } + } + } + } else { + datatype = "int64" + mapper = func(ids []uint64) (interface{}, error) { + switch len(ids) { + case 0: + return nil, nil + case 1: + return int64(ids[0]), nil + default: + return nil, errors.Errorf("BSI field %q has too many values: %v", field.Name, ids) + } + } + } + case FieldTypeDecimal: + datatype = "decimal" + scale := field.Options.Scale + mapper = func(ids []uint64) (_ interface{}, err error) { + switch len(ids) { + case 0: + return nil, nil + case 1: + return pql.NewDecimal(int64(ids[0]), scale), nil + default: + return nil, errors.Errorf("BSI field %q has too many values: %v", field.Name, ids) + } + } + case FieldTypeTimestamp: + datatype = "timestamp" + mapper = func(ids []uint64) (_ interface{}, err error) { + switch len(ids) { + case 0: + return nil, nil + case 1: + return time.Unix(0, int64(ids[0])*int64(featurebase.TimeUnitNanos(field.Options.TimeUnit))).UTC(), nil + default: + return nil, errors.Errorf("BSI field %q has too many values: %v", field.Name, ids) + } + } + default: + return nil, errors.Errorf("field type %q not yet supported", typ) + } + mappers[i] = mapper + fields[i] = featurebase.ExtractedTableField{ + Name: v, + Type: datatype, + } + } + + var translateCol func(uint64) (featurebase.KeyOrID, error) + if idx.Options.Keys { + translateCol = func(id uint64) (featurebase.KeyOrID, error) { + return featurebase.KeyOrID{Keyed: true, Key: idSet[id]}, nil + } + } else { + translateCol = func(id uint64) (featurebase.KeyOrID, error) { + return featurebase.KeyOrID{ID: id}, nil + } + } + + cols := make([]featurebase.ExtractedTableColumn, len(result.Columns)) + colData := make([]interface{}, len(cols)*len(result.Fields)) + for i, col := range result.Columns { + data := colData[i*len(result.Fields) : (i+1)*len(result.Fields) : (i+1)*len(result.Fields)] + for j, rows := range col.Rows { + v, err := mappers[j](rows) + if err != nil { + return nil, errors.Wrap(err, "translating extracted table value") + } + data[j] = v + } + + colTrans, err := translateCol(col.ColumnID) + if err != nil { + return nil, errors.Wrap(err, "translating column ID in extracted table") + } + + cols[i] = featurebase.ExtractedTableColumn{ + Column: colTrans, + Rows: data, + } + } + + return featurebase.ExtractedTable{ + Fields: fields, + Columns: cols, + }, nil + } + + return result, nil +} + +// validateQueryContext returns a query-appropriate error if the context is done. +func validateQueryContext(ctx context.Context) error { + select { + case <-ctx.Done(): + switch err := ctx.Err(); err { + case context.Canceled: + return featurebase.ErrQueryCancelled + case context.DeadlineExceeded: + return featurebase.ErrQueryTimeout + default: + return err + } + default: + return nil + } +} + +type reduceFunc func(ctx context.Context, prev, v interface{}) interface{} + +type mapResponse struct { + node dax.Address + shards []uint64 + + result interface{} + err error +} + +func callArgString(call *pql.Call, key string) string { + value, ok := call.Args[key] + if !ok { + return "" + } + s, _ := value.(string) + return s +} + +type qualifiedOrchestrator struct { + *orchestrator + qual dax.TableQualifier + schemar schemar.Schemar +} + +func newQualifiedOrchestrator(orch *orchestrator, qual dax.TableQualifier, schemar schemar.Schemar) *qualifiedOrchestrator { + return &qualifiedOrchestrator{ + orchestrator: orch, + qual: qual, + schemar: schemar, + } +} + +func (o *qualifiedOrchestrator) Execute(ctx context.Context, index string, q *pql.Query, shards []uint64, opt *featurebase.ExecOptions) (featurebase.QueryResponse, error) { + resp := featurebase.QueryResponse{} + + tkey, err := o.indexToQualifiedTableKey(ctx, index) + if err != nil { + return resp, errors.Wrap(err, "converting index to qualified table key") + } + + return o.orchestrator.Execute(ctx, string(tkey), q, shards, opt) +} + +func (o *qualifiedOrchestrator) indexToQualifiedTableKey(ctx context.Context, index string) (dax.TableKey, error) { + qtid, err := o.schemar.TableID(ctx, o.qual, dax.TableName(index)) + if err != nil { + return "", errors.Wrap(err, "converting index to qualified table id") + } + return qtid.Key(), nil +} diff --git a/dax/queryer/queryer.go b/dax/queryer/queryer.go new file mode 100644 index 000000000..b73f4f3f5 --- /dev/null +++ b/dax/queryer/queryer.go @@ -0,0 +1,304 @@ +// Package queryer provides the core query-related structs. +package queryer + +import ( + "context" + "fmt" + "net/http" + "strings" + "time" + + featurebase "github.com/molecula/featurebase/v3" + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/encoding/proto" + "github.com/molecula/featurebase/v3/errors" + idkmds "github.com/molecula/featurebase/v3/idk/mds" + "github.com/molecula/featurebase/v3/logger" + featurebase_pql "github.com/molecula/featurebase/v3/pql" + fbproto "github.com/molecula/featurebase/v3/proto" + "github.com/molecula/featurebase/v3/server" + "github.com/molecula/featurebase/v3/sql3/parser" + "github.com/molecula/featurebase/v3/sql3/planner" + plannertypes "github.com/molecula/featurebase/v3/sql3/planner/types" + "github.com/molecula/featurebase/v3/stats" +) + +// Queryer represents the query layer in a Molecula implementation. The idea is +// that the externally-facing Molecula API would proxy query requests to a pool +// of "Queryer" nodes, which handle incoming query requests. +type Queryer struct { + orchestrator *orchestrator + + MDS MDS + router Router + + logger logger.Logger +} + +// New returns a new instance of Queryer. +func New(cfg Config) *Queryer { + fbClient, err := featurebase.NewInternalClient("fakehostname:8080", + &http.Client{}, + featurebase.WithSerializer(proto.Serializer{}), + featurebase.WithPathPrefix(dax.ServicePrefixComputer), + ) + if err != nil { + panic(err) // should be impossible + } + + var logr = logger.NopLogger + if cfg.Logger != nil { + logr = cfg.Logger + } + + q := &Queryer{ + MDS: NewNopMDS(), + router: NewNopRouter(), + orchestrator: &orchestrator{ + schema: NewSchemaInfoAPI(cfg.MDS), + trans: NewMDSTranslator(cfg.MDS), + topology: &MDSTopology{mds: cfg.MDS}, + // TODO(jaffee) using default http.Client probably bad... need to set some timeouts. + client: fbClient, + stats: stats.NopStatsClient, + logger: logr, + }, + logger: logr, + } + + if cfg.MDS != nil { + q.MDS = cfg.MDS + } + if cfg.Router != nil { + q.router = cfg.Router + } + + return q +} + +func (q *Queryer) QuerySQL(ctx context.Context, qual dax.TableQualifier, sql string) (*featurebase.SQLResponse, error) { + start := time.Now() + + if len(sql) > 0 && sql[0] == '[' { + return q.parseAndQueryPQL(ctx, qual, sql) + } + ret := &featurebase.SQLResponse{} + + applyExecutionTime := func() { + ret.ExecutionTime = time.Since(start).Microseconds() + } + + applyError := func(e error) { + ret.Error = e.Error() + applyExecutionTime() + } + + st, err := parser.NewParser(strings.NewReader(sql)).ParseStatement() + if err != nil { + applyError(errors.Wrap(err, "parsing sql")) + return ret, nil + } + + // ComputeAPI + capi := NewQualifiedComputeAPI(qual, q.MDS, q.router) + + // SchemaAPI + sapi := NewQualifiedSchemaAPI(qual, q.MDS) + + // Orchestrator + orch := newQualifiedOrchestrator(q.orchestrator, qual, q.MDS) + + // Importer + imp := newBatchImporter(idkmds.NewImporter(q.MDS, nil), qual, q.MDS) + + // TODO(tlt): this obviously doesn't work; we don't have an API here. We + // need a dax-compatible implementation of the SystemAPI (or at least a + // no-op implementation). + sysapi := &featurebase.FeatureBaseSystemAPI{API: nil} + + pl := planner.NewExecutionPlanner(orch, sapi, sysapi, capi, imp, q.orchestrator.logger, sql) + + planOp, err := pl.CompilePlan(ctx, st) + if err != nil { + applyError(errors.Wrap(err, "compiling plan")) + return ret, nil + } + + // Get a query iterator. + iter, err := planOp.Iterator(ctx, nil) + if err != nil { + applyError(errors.Wrap(err, "getting iterator")) + return ret, nil + } + + // Read schema. + columns := planOp.Schema() + schema := featurebase.SQLSchema{ + Fields: make([]*featurebase.SQLField, len(columns)), + } + for i, col := range columns { + schema.Fields[i] = &featurebase.SQLField{ + Name: col.ColumnName, + Type: col.Type.TypeName(), + } + } + + // Read rows. + data := make([][]interface{}, 0) + var currentRow plannertypes.Row + for currentRow, err = iter.Next(ctx); err == nil; currentRow, err = iter.Next(ctx) { + data = append(data, currentRow) + } + if err != nil && err != plannertypes.ErrNoMoreRows { + applyError(errors.Wrap(err, "getting row")) + return ret, nil + } + + ret.Schema = schema + ret.Data = data + applyExecutionTime() + + return ret, nil +} + +func (q *Queryer) parseAndQueryPQL(ctx context.Context, qual dax.TableQualifier, sql string) (*featurebase.SQLResponse, error) { + var i int + for i = 1; sql[i] != ']'; i++ { + if i == len(sql)-1 { + return nil, errors.Errorf("couldn't parse table name out of '%s'", sql) + } + } + table := sql[1:i] + query := sql[i+1:] + fmt.Println("got table/query", table, query) + + return q.QueryPQL(ctx, qual, dax.TableName(table), query) +} + +func (q *Queryer) QueryPQL(ctx context.Context, qual dax.TableQualifier, table dax.TableName, pql string) (*featurebase.SQLResponse, error) { + // Parse the pql into a pql.Query containing []pql.Call. + qry, err := featurebase_pql.NewParser(strings.NewReader(pql)).Parse() + if err != nil { + return nil, errors.Wrap(err, "parsing pql") + } + if len(qry.Calls) != 1 { + return nil, errors.Errorf("must have exactly 1 query, but got: %+v", qry.Calls) + } + + tkey, err := q.indexToQualifiedTableKey(ctx, qual, string(table)) + if err != nil { + return nil, errors.Wrapf(err, "converting index to qualified table key: %s", table) + } + + results, err := q.orchestrator.Execute(ctx, string(tkey), qry, nil, &featurebase.ExecOptions{}) + if err != nil { + return nil, errors.Wrap(err, "orchestrator.Execute") + } + if len(results.Results) != 1 { + return nil, errors.Errorf("expected single result but got %+v", results.Results) + } + + return PQLResultToQueryResult(results.Results[0]) +} + +func PQLResultToQueryResult(pqlResult interface{}) (*featurebase.SQLResponse, error) { + toTabler, err := server.ToTablerWrapper(pqlResult) + if err != nil { + return nil, errors.Wrap(err, "wrapping as type ToTabler") + } + table, err := toTabler.ToTable() + if err != nil { + return nil, errors.Wrap(err, "ToTable") + } + + return tableResponseToQueryResult(table) +} + +func tableResponseToQueryResult(t *fbproto.TableResponse) (*featurebase.SQLResponse, error) { + qr := &featurebase.SQLResponse{ + Schema: featurebase.SQLSchema{Fields: make([]*featurebase.SQLField, len(t.Headers))}, + Data: make([][]interface{}, len(t.Rows)), + } + for i, ci := range t.Headers { + qr.Schema.Fields[i] = &featurebase.SQLField{Name: ci.Name, Type: datatypeToType(ci.Datatype)} + } + + for i, row := range t.Rows { + qr.Data[i] = rowToSliceInterface(t.Headers, row) + } + + return qr, nil +} + +func datatypeToType(ciDatatype string) string { + switch ciDatatype { + case "string": + return parser.FieldTypeString + case "uint64": + return parser.FieldTypeID + case "float64": + // ?? + panic("float64 doesn't have sql3 field type?") + case "int64": + return parser.FieldTypeInt + case "bool": + return parser.FieldTypeBool + case "decimal": + return parser.FieldTypeDecimal + case "timestamp": + return parser.FieldTypeTimestamp + case "[]string": + return parser.FieldTypeStringSet + case "[]uint64": + return parser.FieldTypeIDSet + // TODO []byte?? + default: + panic(fmt.Sprintf("unknown ColumnInfo Datatype: %s", ciDatatype)) + } +} + +func rowToSliceInterface(header []*fbproto.ColumnInfo, row *fbproto.Row) []interface{} { + ret := make([]interface{}, len(row.Columns)) + for i, col := range row.Columns { + switch header[i].Datatype { + case "string": + ret[i] = col.GetStringVal() + case "uint64": + ret[i] = col.GetUint64Val() + case "int64": + ret[i] = col.GetInt64Val() + case "bool": + ret[i] = col.GetBoolVal() + case "[]byte": + ret[i] = col.GetBlobVal() + case "[]uint64": + ret[i] = col.GetUint64ArrayVal() + case "[]string": + ret[i] = col.GetStringArrayVal() + case "float64": + ret[i] = col.GetFloat64Val() + case "decimal": + dec := col.GetDecimalVal() + ret[i] = featurebase_pql.NewDecimal(dec.Value, dec.Scale) + case "timestamp": + ret[i] = col.GetTimestampVal() + default: + panic(fmt.Sprintf("don't know how to get value for columninfo datatype %s, val: %+v, type: %[2]T", header[i].Datatype, col.ColumnVal)) + } + } + return ret +} + +// TODO(tlt): this method was copied from queryer/batchImporter. Can we centralize +// this logic? +func (q *Queryer) indexToQualifiedTableKey(ctx context.Context, qual dax.TableQualifier, index string) (dax.TableKey, error) { + if strings.HasPrefix(index, dax.PrefixTable+dax.TableKeyDelimiter) { + return dax.TableKey(index), nil + } + + qtid, err := q.MDS.TableID(ctx, qual, dax.TableName(index)) + if err != nil { + return "", errors.Wrap(err, "converting index to qualified table id") + } + return qtid.Key(), nil +} diff --git a/dax/queryer/router.go b/dax/queryer/router.go new file mode 100644 index 000000000..2b99cea88 --- /dev/null +++ b/dax/queryer/router.go @@ -0,0 +1,23 @@ +package queryer + +import ( + "github.com/molecula/featurebase/v3/dax" +) + +type Router interface { + Importer(addr dax.Address) Importer +} + +// Ensure type implements interface. +var _ Router = &NopRouter{} + +// NopRouter is a no-op implementation of the Router interface. +type NopRouter struct{} + +func NewNopRouter() *NopRouter { + return &NopRouter{} +} + +func (d *NopRouter) Importer(addr dax.Address) Importer { + return nil +} diff --git a/dax/queryer/schema_api.go b/dax/queryer/schema_api.go new file mode 100644 index 000000000..27e85b82e --- /dev/null +++ b/dax/queryer/schema_api.go @@ -0,0 +1,462 @@ +package queryer + +import ( + "context" + "fmt" + "time" + + pilosa "github.com/molecula/featurebase/v3" + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/mds/schemar" + "github.com/molecula/featurebase/v3/errors" + "github.com/molecula/featurebase/v3/pql" +) + +// Ensure type implements interface. +var _ pilosa.SchemaInfoAPI = (*schemaInfoAPI)(nil) + +type schemaInfoAPI struct { + schemar schemar.Schemar +} + +func NewSchemaInfoAPI(schemar schemar.Schemar) *schemaInfoAPI { + return &schemaInfoAPI{ + schemar: schemar, + } +} + +func (a *schemaInfoAPI) IndexInfo(ctx context.Context, indexName string) (*pilosa.IndexInfo, error) { + qtid := dax.TableKey(indexName).QualifiedTableID() + tbl, err := a.schemar.Table(ctx, qtid) + if err != nil { + return nil, errors.Wrap(err, "getting table for indexinfo") + } + + return daxTableToFeaturebaseIndexInfo(tbl, false) +} + +func (a *schemaInfoAPI) FieldInfo(ctx context.Context, indexName, fieldName string) (*pilosa.FieldInfo, error) { + qtid := dax.TableKey(indexName).QualifiedTableID() + tbl, err := a.schemar.Table(ctx, qtid) + fldName := dax.FieldName(fieldName) + + if err != nil { + return nil, errors.Wrap(err, "getting table for fieldinfo") + } + + fld, ok := tbl.Field(dax.FieldName(fieldName)) + if !ok { + return nil, dax.NewErrFieldDoesNotExist(fldName) + } + + return daxFieldToFeaturebaseFieldInfo(fld) +} + +// daxTableToFeaturebaseIndexInfo converts a dax.Table to a +// featurebase.IndexInfo. If useName is true, the IndexInfo.Name value will +// be set to the qualified table name. Otherwise it will be set to the table key. +func daxTableToFeaturebaseIndexInfo(qtbl *dax.QualifiedTable, useName bool) (*pilosa.IndexInfo, error) { + name := string(qtbl.Key()) + if useName { + name = string(qtbl.Name) + } + ii := &pilosa.IndexInfo{ + Name: name, + CreatedAt: 0, + Options: pilosa.IndexOptions{ + Keys: qtbl.StringKeys(), + TrackExistence: true, + }, + ShardWidth: pilosa.ShardWidth, + } + + // fields + fields := make([]*pilosa.FieldInfo, len(qtbl.Fields)) + var err error + for i := range qtbl.Fields { + fields[i], err = daxFieldToFeaturebaseFieldInfo(qtbl.Fields[i]) + if err != nil { + return nil, errors.Wrap(err, "converting field to FieldInfo") + } + } + ii.Fields = fields + + return ii, nil +} + +// daxFieldToFeaturebaseFieldInfo converts a dax.Field to a +// featurebase.FieldInfo. +func daxFieldToFeaturebaseFieldInfo(field *dax.Field) (*pilosa.FieldInfo, error) { + var timeUnit string + var base int64 + min := field.Options.Min + max := field.Options.Max + + switch field.Type { + case dax.FieldTypeTimestamp: + timestampOptions, err := daxFieldOptionsToFeaturebaseTimestamp(field.Options) + if err != nil { + return nil, errors.Wrap(err, "getting timestamp options") + } + timeUnit = timestampOptions.TimeUnit + base = timestampOptions.Base + min = timestampOptions.Min + max = timestampOptions.Max + } + + fi := &pilosa.FieldInfo{ + Name: string(field.Name), + CreatedAt: 0, // TODO(tlt): we need to handle this on MDS schemar + Options: pilosa.FieldOptions{ + Type: featurebaseFieldType(field), + Base: base, + Min: min, + Max: max, + Scale: field.Options.Scale, + Keys: field.StringKeys(), + NoStandardView: field.Options.NoStandardView, + CacheType: field.Options.CacheType, + CacheSize: field.Options.CacheSize, + TimeUnit: timeUnit, + TimeQuantum: pilosa.TimeQuantum(field.Options.TimeQuantum), + TTL: field.Options.TTL, + ForeignIndex: field.Options.ForeignIndex, + }, + Views: nil, // TODO: do we need views populated? + } + + return fi, nil +} + +// featurebaseFieldType returns the featurebase.FieldType for the given +// dax.Field. +func featurebaseFieldType(f *dax.Field) string { + switch f.Type { + case dax.FieldTypeID, dax.FieldTypeString: + if f.Name == dax.PrimaryKeyFieldName { + return string(f.Type) + } + return "mutex" + case dax.FieldTypeIDSet, dax.FieldTypeStringSet: + if f.Options.TimeQuantum != "" { + return "time" + } + return "set" + default: + return string(f.Type) + } +} + +// featurebaseFieldOptionsToEpoch produces an Epoch (time.Time) value based on +// the given featurebase FieldOptions. +func featurebaseFieldOptionsToEpoch(fo *pilosa.FieldOptions) time.Time { + epochNano := fo.Base * pilosa.TimeUnitNanos(fo.TimeUnit) + return time.Unix(0, epochNano) +} + +// daxFieldOptionsToFeaturebaseTimestamp produces a featurebase.FieldOptions +// value with the applicable options populated. +func daxFieldOptionsToFeaturebaseTimestamp(fo dax.FieldOptions) (*pilosa.FieldOptions, error) { + out := &pilosa.FieldOptions{} + + // Check if the epoch will overflow when converted to nano. + if err := pilosa.CheckEpochOutOfRange(fo.Epoch, pilosa.MinTimestampNano, pilosa.MaxTimestampNano); err != nil { + return nil, errors.Wrap(err, "checking overflow") + } + + out.TimeUnit = fo.TimeUnit + out.Base = fo.Epoch.UnixNano() / pilosa.TimeUnitNanos(fo.TimeUnit) + out.Min = pql.NewDecimal(pilosa.MinTimestamp.UnixNano()/pilosa.TimeUnitNanos(fo.TimeUnit), 0) + out.Max = pql.NewDecimal(pilosa.MaxTimestamp.UnixNano()/pilosa.TimeUnitNanos(fo.TimeUnit), 0) + + return out, nil +} + +func featurebaseFieldOptionSliceToDaxField(name string, opts []pilosa.FieldOption) (*dax.Field, error) { + fo := &pilosa.FieldOptions{} + for _, opt := range opts { + if err := opt(fo); err != nil { + return nil, errors.Wrap(err, "applying field option") + } + } + + return featurebaseFieldOptionsToDaxField(name, fo) +} + +func featurebaseFieldOptionsToDaxField(name string, fo *pilosa.FieldOptions) (*dax.Field, error) { + // Initialize field options; to be overridden based on field type + // specific options. Unless determined otherwise, the defaults for these + // values are applied in sql3/planner/createtable.go, so we don't + // initialize with defaults here. In other words, we set these value to + // exactly as we receive them from the caller. + var fieldType dax.FieldType + var min pql.Decimal + var max pql.Decimal + var scale int64 + var cacheType string + var cacheSize uint32 + var timeUnit string + var epoch time.Time + var foreignIndex string + var timeQuantum dax.TimeQuantum + + switch fo.Type { + case pilosa.FieldTypeMutex: + if fo.Keys { + fieldType = dax.FieldTypeString + } else { + fieldType = dax.FieldTypeID + } + cacheType = fo.CacheType + cacheSize = fo.CacheSize + case pilosa.FieldTypeSet: + if fo.Keys { + fieldType = dax.FieldTypeStringSet + } else { + fieldType = dax.FieldTypeIDSet + } + cacheType = fo.CacheType + cacheSize = fo.CacheSize + case pilosa.FieldTypeInt: + min = fo.Min + max = fo.Max + fieldType = dax.FieldTypeInt + foreignIndex = fo.ForeignIndex + case pilosa.FieldTypeDecimal: + min = fo.Min + max = fo.Max + scale = fo.Scale + fieldType = dax.FieldTypeDecimal + case pilosa.FieldTypeTimestamp: + epoch = featurebaseFieldOptionsToEpoch(fo) + timeUnit = fo.TimeUnit + fieldType = dax.FieldTypeTimestamp + case pilosa.FieldTypeBool: + fieldType = dax.FieldTypeBool + case pilosa.FieldTypeTime: + if fo.Keys { + fieldType = dax.FieldTypeStringSet + } else { + fieldType = dax.FieldTypeIDSet + } + timeQuantum = dax.TimeQuantum(fo.TimeQuantum) + default: + return nil, errors.New(errors.ErrUncoded, fmt.Sprintf("unhandled featurebase field type: %s", fo.Type)) + } + + daxField := &dax.Field{ + Name: dax.FieldName(name), + Type: fieldType, + Options: dax.FieldOptions{ + Min: min, + Max: max, + Scale: scale, + NoStandardView: fo.NoStandardView, + CacheType: cacheType, + CacheSize: cacheSize, + TimeUnit: timeUnit, + Epoch: epoch, + TimeQuantum: timeQuantum, + TTL: fo.TTL, + ForeignIndex: foreignIndex, + }, + } + + return daxField, nil +} + +// Ensure type implements interface. +var _ pilosa.SchemaAPI = (*qualifiedSchemaAPI)(nil) + +// qualifiedSchemaAPI is a wrapper around schemaAPI. It is initialized with a +// TableQualifer, and it uses this qualifer to convert between, for example, +// FeatureBase index name (a string) and TableKey. It requires a Schemar to do +// that lookup/conversion. +type qualifiedSchemaAPI struct { + qual dax.TableQualifier + schemar schemar.Schemar +} + +func NewQualifiedSchemaAPI(qual dax.TableQualifier, schemar schemar.Schemar) *qualifiedSchemaAPI { + return &qualifiedSchemaAPI{ + qual: qual, + schemar: schemar, + } +} + +func (s *qualifiedSchemaAPI) CreateIndexAndFields(ctx context.Context, indexName string, options pilosa.IndexOptions, fields []pilosa.CreateFieldObj) error { + // Make sure the table for this qualifier doesn't already exist. + //_, err := s.schemar.TableID(ctx, s.qual, dax.TableName(indexName)) + _, err := s.schemar.TableID(ctx, s.qual, dax.TableName(indexName)) + // TODO(tlt): the following doesn't work when the TableID() call is made + // over http because the error that comes back is `status code: 400: table + // name 'tbl' does not exist\n`, which does not match the check. We really + // need to be able to check these error codes both directly on the error AND + // when they come back via http. + // if !errors.Is(err, dax.ErrTableNameDoesNotExist) { + // if err != nil { + // return errors.Wrapf(err, "checking if table name already exists: %s, %s", s.qual, indexName) + // } + // return dax.NewErrTableNameExists(dax.TableName(indexName)) + // } + if err == nil { + return dax.NewErrTableNameExists(dax.TableName(indexName)) + } + + partitionN := dax.DefaultPartitionN + // TODO(tlt): until we can thread partitionN through featurebase correctly + // (instead of having it use holder.partitionN or DefaultPartitionN), then + // we can't use a custom partitionN. + // if options.PartitionN > 0 { + // partitionN = options.PartitionN + // } + + // Initialize the fields slice with one additional slot for the primary key + // field. + daxFields := make([]*dax.Field, 0, len(fields)+1) + + // Add the primary key field. + var fieldType dax.FieldType + if options.Keys { + fieldType = dax.FieldTypeString + } else { + fieldType = dax.FieldTypeID + } + daxFields = append(daxFields, &dax.Field{ + Name: dax.PrimaryKeyFieldName, + Type: fieldType, + }) + + // Add the fields provided in the method call. + for _, fldObj := range fields { + daxField, err := featurebaseFieldOptionSliceToDaxField(fldObj.Name, fldObj.Options) + if err != nil { + return errors.Wrap(err, "converting featurebase field options to dax field") + } + + daxFields = append(daxFields, daxField) + } + + tbl := &dax.Table{ + Name: dax.TableName(indexName), + Fields: daxFields, + PartitionN: partitionN, + } + + qtbl := dax.NewQualifiedTable( + s.qual, + tbl, + ) + + return s.schemar.CreateTable(ctx, qtbl) +} + +func (s *qualifiedSchemaAPI) CreateField(ctx context.Context, indexName string, fieldName string, opts ...pilosa.FieldOption) (*pilosa.Field, error) { + tkey, err := s.indexToQualifiedTableKey(ctx, indexName) + if err != nil { + return nil, errors.Wrap(err, "converting index to qualified table key") + } + + daxField, err := featurebaseFieldOptionSliceToDaxField(fieldName, opts) + if err != nil { + return nil, errors.Wrap(err, "converting featurebase field options to dax field") + } + + qtid := tkey.QualifiedTableID() + + if err := s.schemar.CreateField(ctx, qtid, daxField); err != nil { + return nil, errors.New(errors.ErrUncoded, err.Error()) + } + + return nil, nil +} + +func (s *qualifiedSchemaAPI) DeleteField(ctx context.Context, indexName string, fieldName string) error { + tkey, err := s.indexToQualifiedTableKey(ctx, indexName) + if err != nil { + return errors.Wrap(err, "converting index to qualified table key") + } + + qtid := tkey.QualifiedTableID() + fldName := dax.FieldName(fieldName) + + if err := s.schemar.DropField(ctx, qtid, fldName); err != nil { + return errors.New(errors.ErrUncoded, err.Error()) + } + + return nil +} + +func (s *qualifiedSchemaAPI) DeleteIndex(ctx context.Context, indexName string) error { + tkey, err := s.indexToQualifiedTableKey(ctx, indexName) + if err != nil { + return errors.Wrap(err, "converting index to qualified table key") + } + qtid := tkey.QualifiedTableID() + return s.schemar.DropTable(ctx, qtid) +} + +func (s *qualifiedSchemaAPI) IndexInfo(ctx context.Context, indexName string) (*pilosa.IndexInfo, error) { + tkey, err := s.indexToQualifiedTableKey(ctx, indexName) + if err != nil { + return nil, errors.Wrap(err, "converting index to qualified table key") + } + + qtid := tkey.QualifiedTableID() + tbl, err := s.schemar.Table(ctx, qtid) + if err != nil { + return nil, errors.Wrap(err, "getting table for qualified indexinfo") + } + + return daxTableToFeaturebaseIndexInfo(tbl, true) +} + +func (s *qualifiedSchemaAPI) FieldInfo(ctx context.Context, indexName, fieldName string) (*pilosa.FieldInfo, error) { + tkey, err := s.indexToQualifiedTableKey(ctx, indexName) + if err != nil { + return nil, errors.Wrap(err, "converting index to qualified table key") + } + + qtid := tkey.QualifiedTableID() + tbl, err := s.schemar.Table(ctx, qtid) + fldName := dax.FieldName(fieldName) + + if err != nil { + return nil, errors.Wrap(err, "getting table for qualified fieldinfo") + } + + fld, ok := tbl.Field(dax.FieldName(fieldName)) + if !ok { + return nil, dax.NewErrFieldDoesNotExist(fldName) + } + + return daxFieldToFeaturebaseFieldInfo(fld) +} + +func (s *qualifiedSchemaAPI) Schema(ctx context.Context, withViews bool) ([]*pilosa.IndexInfo, error) { + tbls, err := s.schemar.Tables(ctx, s.qual) + if err != nil { + return nil, errors.Wrap(err, "getting tables for qualified schema") + } + + indexes := make([]*pilosa.IndexInfo, len(tbls)) + for i := range tbls { + // This method appears to be used primarily in "SHOW TABLES", and in + // that case we want to return the human friendly table name used when + // creating the table (i.e. not the table key). + indexes[i], err = daxTableToFeaturebaseIndexInfo(tbls[i], true) + if err != nil { + return nil, errors.Wrap(err, "converting table to IndexInfo") + } + } + + return indexes, nil +} + +func (s *qualifiedSchemaAPI) indexToQualifiedTableKey(ctx context.Context, index string) (dax.TableKey, error) { + qtid, err := s.schemar.TableID(ctx, s.qual, dax.TableName(index)) + if err != nil { + return "", errors.Wrap(err, "converting index to qualified table id") + } + return qtid.Key(), nil +} diff --git a/dax/queryer/translator.go b/dax/queryer/translator.go new file mode 100644 index 000000000..ba96600ea --- /dev/null +++ b/dax/queryer/translator.go @@ -0,0 +1,325 @@ +package queryer + +import ( + "context" + "net/http" + + pilosa "github.com/molecula/featurebase/v3" + featurebase_client "github.com/molecula/featurebase/v3/client" + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/mds/controller/partitioner" + "github.com/molecula/featurebase/v3/disco" + "github.com/molecula/featurebase/v3/encoding/proto" + "github.com/molecula/featurebase/v3/errors" +) + +// Ensure type implements interface. +var _ Translator = (*MDSTranslator)(nil) + +type MDSTranslator struct { + mds MDS +} + +func NewMDSTranslator(mds MDS) *MDSTranslator { + return &MDSTranslator{ + mds: mds, + } +} + +func fbClient(address dax.Address) (*featurebase_client.Client, error) { + // Set up a FeatureBase client with address. + return featurebase_client.NewClient(address.String(), + featurebase_client.OptClientRetries(2), + featurebase_client.OptClientTotalPoolSize(1000), + featurebase_client.OptClientPoolSizePerRoute(400), + featurebase_client.OptClientPathPrefix(dax.ServicePrefixComputer), + //featurebase_client.OptClientStatsClient(m.stats), + ) +} + +func (m *MDSTranslator) CreateIndexKeys(ctx context.Context, table string, keys []string) (map[string]uint64, error) { + tkey := dax.TableKey(table) + qtid := tkey.QualifiedTableID() + + qtbl, err := m.mds.Table(ctx, qtid) + if err != nil { + return nil, errors.Wrap(err, "getting table") + } + + partitioner := partitioner.NewPartitioner() + + // Get the partitions (and therefore, nodes) responsible for the keys. + pMap := partitioner.PartitionsForKeys(tkey, qtbl.PartitionN, keys...) + + out := make(map[string]uint64) + for pNum := range pMap { + address, err := m.mds.IngestPartition(ctx, qtid, pNum) + if err != nil { + return nil, errors.Wrapf(err, "calling ingest-partition on table: %s, partition: %d", table, pNum) + } + + fbClient, err := fbClient(address) + if err != nil { + return nil, errors.Wrap(err, "getting featurebase client") + } + + idx := featurebase_client.NewIndex(table) + + m, err := fbClient.CreateIndexKeys(idx, pMap[pNum]...) + if err != nil { + return nil, errors.Wrapf(err, "creating index keys on index: %s, partition: %d", table, pNum) + } + + for k, v := range m { + out[k] = v + } + } + + return out, nil +} + +func (m *MDSTranslator) CreateFieldKeys(ctx context.Context, table string, field string, keys []string) (map[string]uint64, error) { + qtid := dax.TableKey(table).QualifiedTableID() + address, err := m.mds.IngestPartition(ctx, qtid, dax.PartitionNum(0)) + if err != nil { + return nil, errors.Wrapf(err, "calling ingest-partition on table: %s, partition: %d", table, dax.PartitionNum(0)) + } + + fbClient, err := fbClient(address) + if err != nil { + return nil, errors.Wrap(err, "getting featurebase client") + } + + idx := featurebase_client.NewIndex(table) + fld := idx.Field(field) + + return fbClient.CreateFieldKeys(fld, keys...) +} + +func (m *MDSTranslator) FindIndexKeys(ctx context.Context, table string, keys []string) (map[string]uint64, error) { + tkey := dax.TableKey(table) + qtid := tkey.QualifiedTableID() + + qtbl, err := m.mds.Table(ctx, qtid) + if err != nil { + return nil, errors.Wrap(err, "getting table") + } + + partitioner := partitioner.NewPartitioner() + + // Get the partitions (and therefore, nodes) responsible for the keys. + pMap := partitioner.PartitionsForKeys(tkey, qtbl.PartitionN, keys...) + + pNums := make([]dax.PartitionNum, 0, len(pMap)) + for k := range pMap { + pNums = append(pNums, k) + } + + translateNodes, err := m.mds.TranslateNodes(ctx, qtid, pNums...) + if err != nil { + return nil, errors.Wrapf(err, "getting translate nodes for partitions on table: %s", table) + } + + out := make(map[string]uint64) + for _, tnode := range translateNodes { + address := tnode.Address + + fbClient, err := fbClient(address) + if err != nil { + return nil, errors.Wrap(err, "getting featurebase client") + } + + idx := featurebase_client.NewIndex(table) + + nodeKeys := []string{} + for _, pNum := range tnode.Partitions { + nodeKeys = append(nodeKeys, pMap[pNum]...) + } + + m, err := fbClient.FindIndexKeys(idx, nodeKeys...) + if err != nil { + return nil, errors.Wrapf(err, "finding index keys on index: %s", table) + } + + for k, v := range m { + out[k] = v + } + } + + return out, nil +} + +func (m *MDSTranslator) FindFieldKeys(ctx context.Context, table, field string, keys []string) (map[string]uint64, error) { + qtid := dax.TableKey(table).QualifiedTableID() + address, err := m.mds.IngestPartition(ctx, qtid, dax.PartitionNum(0)) + if err != nil { + return nil, errors.Wrapf(err, "calling ingest-partition on table: %s, partition: %d", table, dax.PartitionNum(0)) + } + + fbClient, err := fbClient(address) + if err != nil { + return nil, errors.Wrap(err, "getting featurebase client") + } + + idx := featurebase_client.NewIndex(table) + fld := idx.Field(field) + + return fbClient.FindFieldKeys(fld, keys...) +} + +func (m *MDSTranslator) TranslateIndexIDs(ctx context.Context, index string, ids []uint64) ([]string, error) { + idsByPartition := splitIDsByPartition(index, ids, 1<<20) // TODO(jaffee), don't hardcode shardwidth...need to get this from index info + daxPartitions := make([]dax.PartitionNum, 0) + for partition := range idsByPartition { + daxPartitions = append(daxPartitions, partition) + } + + qtid := dax.TableKey(index).QualifiedTableID() + + nodes, err := m.mds.TranslateNodes(ctx, qtid, daxPartitions...) + if err != nil { + return nil, errors.Wrapf(err, "calling translate-nodes on table: %s, partitions: %v", index, daxPartitions) + } + + // get translation from each node + idToKey := make(map[uint64]string) + for _, node := range nodes { + fbClient, err := fbClient(node.Address) + if err != nil { + return nil, errors.Wrap(err, "getting featurebase client") + } + + reqIDs := make([]uint64, 0) + for _, partition := range node.Partitions { + reqIDs = append(reqIDs, idsByPartition[partition]...) + } + + strings, err := makeTranslateIDsRequest(fbClient, index, "", reqIDs) + if err != nil { + return nil, errors.Wrapf(err, "translating on %v", node.Address) + } + for i, s := range strings { + idToKey[reqIDs[i]] = s + } + } + + ret := make([]string, len(ids)) + for i, id := range ids { + ret[i] = idToKey[id] + } + return ret, nil +} + +func (m *MDSTranslator) TranslateIndexIDSet(ctx context.Context, table string, ids map[uint64]struct{}) (map[uint64]string, error) { + idList := make([]uint64, 0, len(ids)) + for id := range ids { + idList = append(idList, id) + } + + stringList, err := m.TranslateIndexIDs(ctx, table, idList) + if err != nil { + return nil, errors.Wrapf(err, "translating index ids on table: %s", table) + } + + ret := make(map[uint64]string) + for i, id := range idList { + ret[id] = stringList[i] + } + return ret, nil +} +func (m *MDSTranslator) TranslateFieldIDs(ctx context.Context, table, field string, ids map[uint64]struct{}) (map[uint64]string, error) { + idList := make([]uint64, 0, len(ids)) + for id := range ids { + idList = append(idList, id) + } + + stringList, err := m.TranslateFieldListIDs(ctx, table, field, idList) + if err != nil { + return nil, errors.Wrapf(err, "translating field ids on field: %s, %s", table, field) + } + + ret := make(map[uint64]string) + for i, id := range idList { + ret[id] = stringList[i] + } + return ret, nil +} +func (m *MDSTranslator) TranslateFieldListIDs(ctx context.Context, index, field string, ids []uint64) ([]string, error) { + qtid := dax.TableKey(index).QualifiedTableID() + address, err := m.mds.IngestPartition(ctx, qtid, dax.PartitionNum(0)) + if err != nil { + return nil, errors.Wrapf(err, "calling ingest-partition on table: %s, partition: %d", index, dax.PartitionNum(0)) + } + + fbClient, err := fbClient(address) + if err != nil { + return nil, errors.Wrap(err, "getting featurebase client") + } + + return makeTranslateIDsRequest(fbClient, index, field, ids) +} + +func makeTranslateIDsRequest(fbClient *featurebase_client.Client, table, field string, ids []uint64) ([]string, error) { + method := "POST" + path := "/" + dax.ServicePrefixComputer + "/internal/translate/ids" + headers := map[string]string{ + "Content-Type": "application/x-protobuf", + "Accept": "application/x-protobuf", + } + + req := &pilosa.TranslateIDsRequest{ + Index: table, + Field: field, + IDs: ids, + } + + ser := proto.Serializer{} + + data, err := ser.Marshal(req) + if err != nil { + return nil, errors.Wrap(err, "marshaling translate ids request") + } + + status, body, err := fbClient.HTTPRequest(method, path, data, headers) + if err != nil { + return nil, errors.Wrap(err, "http request") + } else if status != http.StatusOK { + return nil, errors.Wrapf(err, "http request status code: %d", status) + } + + idsResp := &pilosa.TranslateIDsResponse{} + + if err := ser.Unmarshal(body, idsResp); err != nil { + return nil, errors.Wrap(err, "unmarshaling translate ids request") + } + + return idsResp.Keys, nil +} + +func splitIDsByShard(ids []uint64, shardWidth uint64) map[dax.ShardNum][]uint64 { + ret := make(map[dax.ShardNum][]uint64) + for _, id := range ids { + shardIDs, ok := ret[dax.ShardNum(id/shardWidth)] + if !ok { + shardIDs = make([]uint64, 0) + } + ret[dax.ShardNum(id/shardWidth)] = append(shardIDs, id) + } + return ret +} + +func splitIDsByPartition(index string, ids []uint64, shardWidth uint64) map[dax.PartitionNum][]uint64 { + idsByShard := splitIDsByShard(ids, shardWidth) + + partitioner := partitioner.NewPartitioner() + + ret := make(map[dax.PartitionNum][]uint64) + for shard, ids := range idsByShard { + // get partition for shard + // TODO: need to get partitionN from the table, instead of using the default. + partitionNum := partitioner.ShardToPartition(dax.TableKey(index), shard, disco.DefaultPartitionN) + + ret[partitionNum] = append(ret[partitionNum], ids...) + } + return ret +} diff --git a/dax/role.go b/dax/role.go new file mode 100644 index 000000000..db1e2ee93 --- /dev/null +++ b/dax/role.go @@ -0,0 +1,57 @@ +package dax + +// RoleType represents a role type which a worker node can act as. +type RoleType string + +const ( + RoleTypeCompute RoleType = "compute" + RoleTypeTranslate RoleType = "translate" +) + +// RoleTypes is a list of RoleType, used primarily to introduce helper methods +// on the list. +type RoleTypes []RoleType + +// Contains returns true if the given RoleType is in RoleTypes. +func (rt RoleTypes) Contains(t RoleType) bool { + for i := range rt { + if rt[i] == t { + return true + } + } + return false +} + +// Role is an interface for any role which a worker node can assume. +type Role interface { + Type() RoleType +} + +// Ensure type implements interface. +var _ Role = &ComputeRole{} +var _ Role = &TranslateRole{} + +// ComputeRole is a role specific to compute nodes. +type ComputeRole struct { + TableKey TableKey `json:"table-key"` + Shards Shards `json:"shards"` +} + +// Type returns the type for ComputeRole. This is mainly to impolement the Role +// interface. +func (cr *ComputeRole) Type() RoleType { + return RoleTypeCompute +} + +// TranslateRole is a role specific to translate nodes. +type TranslateRole struct { + TableKey TableKey `json:"table-key"` + Partitions Partitions `json:"partitions"` + Fields FieldVersions `json:"fields"` +} + +// Type returns the type for TransteRole. This is mainly to impolement the Role +// interface. +func (cr *TranslateRole) Type() RoleType { + return RoleTypeTranslate +} diff --git a/dax/server/config.go b/dax/server/config.go new file mode 100644 index 000000000..3419997fe --- /dev/null +++ b/dax/server/config.go @@ -0,0 +1,303 @@ +// Copyright 2021 Molecula Corp. All rights reserved. +package server + +import ( + "context" + "fmt" + "log" + "net" + "strconv" + "strings" + "time" + + "github.com/molecula/featurebase/v3/dax/queryer" + "github.com/molecula/featurebase/v3/dax/snapshotter" + "github.com/molecula/featurebase/v3/dax/writelogger" + "github.com/molecula/featurebase/v3/errors" + fbserver "github.com/molecula/featurebase/v3/server" +) + +const ( + defaultBindPort = "8080" + defaultStorageMethod = "boltdb" +) + +// Config represents the configuration for the command. +type Config struct { + // Bind is the host:port on which Pilosa will listen. + Bind string `toml:"bind"` + + // Advertise is the address advertised by the server to other nodes + // in the cluster. It should be reachable by all other nodes and should + // route to an interface that Bind is listening on. + Advertise string `toml:"advertise"` + + // Verbose toggles verbose logging which can be useful for debugging. + Verbose bool `toml:"verbose"` + + // LogPath configures where Pilosa will write logs. + LogPath string `toml:"log-path"` + + MDS MDSOptions `toml:"mds"` + WriteLogger WriteLoggerOptions `toml:"writelogger"` + Snapshotter SnapshotterOptions `toml:"snapshotter"` + Queryer QueryerOptions `toml:"queryer"` + Computer ComputerOptions `toml:"computer"` + + // Storage methods. + StorageMethod string `toml:"storage-method"` + StorageDSN string `toml:"storage-dsn"` +} + +type MDSOptions struct { + Run bool `toml:"run"` + Config MDSConfig `toml:"config"` +} + +type MDSConfig struct { + RegistrationBatchTimeout time.Duration `toml:"registration-batch-timeout"` +} + +type WriteLoggerOptions struct { + Run bool `toml:"run"` + Config writelogger.Config `toml:"config"` +} + +type SnapshotterOptions struct { + Run bool `toml:"run"` + Config snapshotter.Config `toml:"config"` +} + +type QueryerOptions struct { + Run bool `toml:"run"` + Config queryer.Config `toml:"config"` +} + +type ComputerOptions struct { + Run bool `toml:"run"` + Config fbserver.Config `toml:"config"` +} + +// NewConfig returns an instance of Config with default options. +func NewConfig() *Config { + c := &Config{ + MDS: MDSOptions{ + Config: MDSConfig{ + RegistrationBatchTimeout: time.Second * 3, + }, + }, + Bind: ":" + defaultBindPort, + Computer: ComputerOptions{ + Config: *fbserver.NewConfig(), + }, + StorageMethod: defaultStorageMethod, + } + return c +} + +// MustValidate is carried over from server/config.go; it's just a stubbed out +// no-op for now. +func (c *Config) MustValidate() { + err := c.validate() + if err != nil { + panic(err) + } +} + +func (c *Config) validate() error { + return nil +} + +// validateAddrs controls the address fields in the Config object +// and fills in any blanks. +// The addresses fields must be guaranteed by the caller to either be +// completely empty, or have both a host part and a port part +// separated by a colon. In the latter case either can be empty to +// indicate it's left unspecified. +func (c *Config) validateAddrs(ctx context.Context) error { + // Validate the advertise address. + advScheme, advHost, advPort, err := validateAdvertiseAddr(ctx, c.Advertise, c.Bind, defaultBindPort) + if err != nil { + return errors.Wrapf(err, "validating advertise address") + } + c.Advertise = schemeHostPortString(advScheme, advHost, advPort) + + // Validate the listen address. + listenScheme, listenHost, listenPort, err := validateListenAddr(ctx, c.Bind, defaultBindPort) + if err != nil { + return errors.Wrap(err, "validating listen address") + } + c.Bind = schemeHostPortString(listenScheme, listenHost, listenPort) + + return nil +} + +// validateAdvertiseAddr validates and normalizes an address accessible +// Ensures that if the "host" part is empty, it gets filled in with +// the configured listen address if any, otherwise it makes a best +// guess at the outbound IP address. +// Returns scheme, host, port as strings. +func validateAdvertiseAddr(ctx context.Context, advAddr, listenAddr, defaultPort string) (string, string, string, error) { + listenScheme, listenHost, listenPort, err := splitAddr(listenAddr, defaultPort) + if err != nil { + return "", "", "", errors.Wrap(err, "getting listen address") + } + + advScheme, advHostPort := splitScheme(advAddr) + advHost, advPort := "", "" + if advHostPort != "" { + var err error + advHost, advPort, err = net.SplitHostPort(advHostPort) + if err != nil { + return "", "", "", errors.Wrapf(err, "splitting host port: %s", advHostPort) + } + } + // If no advertise scheme was specified, use the one from + // the listen address. + if advScheme == "" { + advScheme = listenScheme + } + // If there was no port number, reuse the one from the listen + // address. + if advPort == "" || advPort == "0" { + advPort = listenPort + } + // Resolve non-numeric to numeric. + portNumber, err := net.DefaultResolver.LookupPort(ctx, "tcp", advPort) + if err != nil { + return "", "", "", errors.Wrapf(err, "looking up non-numeric port: %v", advPort) + } + advPort = strconv.Itoa(portNumber) + + // If the advertise host is empty, then we have two cases. + if advHost == "" { + if listenHost == "0.0.0.0" { + advHost = outboundIP().String() + } else { + advHost = listenHost + } + } + return advScheme, advHost, advPort, nil +} + +// validateListenAddr validates and normalizes an address suitable for +// use with net.Listen(). This accepts an empty "host" part to signify +// the default (localhost) should be used. Rresolves host names to IP +// addresses. +// Returns scheme, host, port as strings. +func validateListenAddr(ctx context.Context, addr, defaultPort string) (string, string, string, error) { + scheme, host, port, err := splitAddr(addr, defaultPort) + if err != nil { + return "", "", "", errors.Wrap(err, "getting listen address") + } + rHost, rPort, err := resolveAddr(ctx, host, port) + if err != nil { + return "", "", "", errors.Wrap(err, "resolving address") + } + return scheme, rHost, rPort, nil +} + +func schemeHostPortString(scheme, host, port string) string { + var s string + if scheme != "" { + s += fmt.Sprintf("%s://", scheme) + } + return s + net.JoinHostPort(host, port) +} + +// splitAddr returns scheme, host, port as strings. +func splitAddr(addr string, defaultPort string) (string, string, string, error) { + scheme, hostPort := splitScheme(addr) + host, port := "", "" + if hostPort != "" { + var err error + host, port, err = net.SplitHostPort(hostPort) + if err != nil { + return "", "", "", errors.Wrapf(err, "splitting host port: %s", hostPort) + } + } + // It's not ideal to have a default here, but the alterative + // results in a port of 0, which causes Pilosa to listen on + // a random port. + if port == "" { + port = defaultPort + } + return scheme, host, port, nil +} + +// resolveAddr resolves non-numeric addresses to numeric (IP, port) addresses. +func resolveAddr(ctx context.Context, host, port string) (string, string, error) { + resolver := net.DefaultResolver + + // Resolve the port number. This may translate service names + // e.g. "postgresql" to a numeric value. + portNumber, err := resolver.LookupPort(ctx, "tcp", port) + if err != nil { + return "", "", errors.Wrapf(err, "resolving up port: %v", port) + } + port = strconv.Itoa(portNumber) + + // Resolve the address. + if host == "" || host == "localhost" { + return host, port, nil + } + + addr, err := lookupAddr(ctx, resolver, host) + if err != nil { + return "", "", errors.Wrap(err, "looking up address") + } + return addr, port, nil +} + +// splitScheme returns two strings: the scheme and the hostPort. +func splitScheme(addr string) (string, string) { + parts := strings.SplitN(addr, "://", 2) + if len(parts) == 1 { + return "", addr + } + return parts[0], parts[1] +} + +// lookupAddr resolves the given address/host to an IP address. If +// multiple addresses are resolved, it returns the first IPv4 address +// available if there is one, otherwise the first address. +func lookupAddr(ctx context.Context, resolver *net.Resolver, host string) (string, error) { + // Resolve the IP address or hostname to an IP address. + addrs, err := resolver.LookupIPAddr(ctx, host) + if err != nil { + return "", errors.Wrap(err, "looking up IP addresses") + } + if len(addrs) == 0 { + return "", fmt.Errorf("cannot resolve %q to an address", host) + } + + // LookupIPAddr() can return a mix of IPv6 and IPv4 + // addresses. Return the first IPv4 address if possible. + for _, addr := range addrs { + if ip := addr.IP.To4(); ip != nil { + return ip.String(), nil + } + } + + // No IPv4 address, return the first resolved address instead. + return addrs[0].String(), nil +} + +// outboundIP gets the preferred outbound ip of this machine. +func outboundIP() net.IP { + // This is not actually making a connection to 8.8.8.8. + // net.Dial() selects the IP address that would be used + // if an actual connection to 8.8.8.8 were made, so this + // choice of address is just meant to ensure that an + // external address is returned (as opposed to a local + // address like 127.0.0.1). + conn, err := net.Dial("udp", "8.8.8.8:80") + if err != nil { + log.Fatal(err) + } + defer conn.Close() + + localAddr := conn.LocalAddr().(*net.UDPAddr) + + return localAddr.IP +} diff --git a/dax/server/dup.go b/dax/server/dup.go new file mode 100644 index 000000000..7c1fe6521 --- /dev/null +++ b/dax/server/dup.go @@ -0,0 +1,15 @@ +// Copyright 2022 Molecula Corp. All rights reserved. +//go:build darwin || (linux && !arm64) +// +build darwin linux,!arm64 + +package server + +import ( + "syscall" +) + +// dup is an alias for syscall.Dup2 on darwin-amd64, darwin-arm64, +// linux-amd64, linux-arm or syscall.Dup3 on linux-arm64 +func (m *Command) dup(oldfd int, newfd int) error { + return syscall.Dup2(oldfd, newfd) +} diff --git a/dax/server/dup_arm64.go b/dax/server/dup_arm64.go new file mode 100644 index 000000000..0338e2d15 --- /dev/null +++ b/dax/server/dup_arm64.go @@ -0,0 +1,15 @@ +// Copyright 2022 Molecula Corp. All rights reserved. +//go:build linux && arm64 +// +build linux,arm64 + +package server + +import ( + "syscall" +) + +// dup is an alias for syscall.Dup2 on darwin-amd64, darwin-arm64, +// linux-amd64, linux-arm or syscall.Dup3 on linux-arm64 +func (m *Command) dup(oldfd int, newfd int) error { + return syscall.Dup3(oldfd, newfd, 0) +} diff --git a/dax/server/server.go b/dax/server/server.go new file mode 100644 index 000000000..7de324fb4 --- /dev/null +++ b/dax/server/server.go @@ -0,0 +1,523 @@ +// Copyright 2021 Molecula Corp. All rights reserved. +// +// Package server contains the `pilosa server` subcommand which runs Pilosa +// itself. The purpose of this package is to define an easily tested Command +// object which handles interpreting configuration and setting up all the +// objects that Pilosa needs. + +package server + +import ( + "context" + "crypto/tls" + "encoding/json" + "io" + "math/rand" + "net" + "os" + "os/signal" + "syscall" + "time" + + "golang.org/x/sync/errgroup" + + featurebase "github.com/molecula/featurebase/v3" + "github.com/molecula/featurebase/v3/dax" + daxhttp "github.com/molecula/featurebase/v3/dax/http" + "github.com/molecula/featurebase/v3/dax/mds" + mdsclient "github.com/molecula/featurebase/v3/dax/mds/client" + controlleralpha "github.com/molecula/featurebase/v3/dax/mds/controller/alpha" + controllerhttp "github.com/molecula/featurebase/v3/dax/mds/controller/http" + "github.com/molecula/featurebase/v3/dax/queryer" + queryeralpha "github.com/molecula/featurebase/v3/dax/queryer/alpha" + "github.com/molecula/featurebase/v3/dax/snapshotter" + snapshotterclient "github.com/molecula/featurebase/v3/dax/snapshotter/client" + "github.com/molecula/featurebase/v3/dax/writelogger" + writeloggerclient "github.com/molecula/featurebase/v3/dax/writelogger/client" + "github.com/molecula/featurebase/v3/errors" + "github.com/molecula/featurebase/v3/logger" + fbnet "github.com/molecula/featurebase/v3/net" + featurebaseserver "github.com/molecula/featurebase/v3/server" +) + +// Command represents the state of the dax server command. +type Command struct { + // Server *pilosa.Server + + // Configuration. + Config *Config + + Handler featurebase.HandlerI + + // done will be closed when Command.Close() is called + done chan struct{} + + // Standard input/output + *featurebase.CmdIO + + ln net.Listener + listenURI *fbnet.URI + advertiseURI *fbnet.URI + tlsConfig *tls.Config + + // registerFns is a list of functions to call once the service is up and + // running. This is typically used to register the service with MDS. + registerFns []registerFn + + // checkInFn is a function to call periodically in order to check-in with a + // monitorinig service such as MDS. + checkInFn checkInFn + + logger logger.Logger + logOutput io.Writer +} + +type registerFn func() error +type checkInFn func() error + +type CommandOption func(c *Command) error + +func OptCommandConfig(config *Config) CommandOption { + return func(c *Command) error { + defer c.Config.MustValidate() + if c.Config != nil { + // c.Config.Etcd = config.Etcd + // c.Config.Auth = config.Auth + // c.Config.TLS = config.TLS + // c.Config.Controller = config.Controller + // c.Config.WriteLogger = config.WriteLogger + return nil + } + c.Config = config + return nil + } +} + +// NewCommand returns a new instance of Command. +func NewCommand(stdin io.Reader, stdout, stderr io.Writer, opts ...CommandOption) *Command { + c := &Command{ + Config: NewConfig(), + + CmdIO: featurebase.NewCmdIO(stdin, stdout, stderr), + + registerFns: make([]registerFn, 0), + + done: make(chan struct{}), + } + + for _, opt := range opts { + err := opt(c) + if err != nil { + panic(err) + // TODO: Return error instead of panic? + } + } + + return c +} + +// Start starts the DAX server. +func (m *Command) Start() (err error) { + // Seed random number generator + rand.Seed(time.Now().UTC().UnixNano()) + + if err := m.setupServer(); err != nil { + return errors.Wrap(err, "setting up server") + } + + // // Initialize server. + // if err = m.Server.Open(); err != nil { + // return errors.Wrap(err, "opening server") + // } + + // Serve HTTP. + go func() { + if err := m.Handler.Serve(); err != nil { + m.logger.Errorf("handler serve error (dax): %v", err) + } + }() + m.logger.Printf("listening as %s\n", m.listenURI) + + // Register the service(s) by calling any registerFn they have implemented. + for i := range m.registerFns { + if err := m.registerFns[i](); err != nil { + return errors.Wrap(err, "calling register function") + } + } + + // Start the "check-in" background process which periodically checks in with + // MDS. + go m.checkIn() + + return nil +} + +// checkIn calls the CheckIn function set on m.checkInFn every interval period. +// If the interval period is 0, the check-in is disabled. +func (m *Command) checkIn() { + interval := m.Config.Computer.Config.CheckInInterval + + if interval == 0 || m.checkInFn == nil { + return + } + + for { + select { + case <-m.done: + return + case <-time.After(interval): + m.logger.Debugf("node check-in in last %s, address: %s", interval, m.Config.Advertise) + if err := m.checkInFn(); err != nil { + m.logger.Errorf("checking in node: %s, %v", m.Config.Advertise, err) + } + } + } +} + +// Wait waits for the server to be closed or interrupted. +func (m *Command) Wait() error { + // First SIGKILL causes server to shut down gracefully. + c := make(chan os.Signal, 2) + signal.Notify(c, os.Interrupt, syscall.SIGTERM) + select { + case sig := <-c: + m.logger.Infof("received signal '%s', gracefully shutting down...\n", sig.String()) + + // Second signal causes a hard shutdown. + go func() { <-c; os.Exit(1) }() + return errors.Wrap(m.Close(), "closing command") + case <-m.done: + m.logger.Infof("server closed externally") + return nil + } +} + +// Close shuts down the server. +func (m *Command) Close() error { + select { + case <-m.done: + return nil + default: + eg := errgroup.Group{} + //eg.Go(m.Server.Close) + + err := eg.Wait() + //_ = testhook.Closed(pilosa.NewAuditor(), m, nil) + close(m.done) + + return errors.Wrap(err, "closing everything") + } +} + +// // ParseConfig parses s into a Config. +// func ParseConfig(s string) (Config, error) { +// var c Config +// err := toml.Unmarshal([]byte(s), &c) +// return c, err +// } + +// // expandDirName was copied from pilosa/server.go. +// // TODO: consider centralizing this if we need this across packages. +// func expandDirName(path string) (string, error) { +// prefix := "~" + string(filepath.Separator) +// if strings.HasPrefix(path, prefix) { +// HomeDir := os.Getenv("HOME") +// if HomeDir == "" { +// return "", errors.New("data directory not specified and no home dir available") +// } +// return filepath.Join(HomeDir, strings.TrimPrefix(path, prefix)), nil +// } +// return path, nil +// } + +// setupServer uses the configuration to set up this server. +func (m *Command) setupServer() error { + // Set up logger. + if err := m.setupLogger(); err != nil { + return errors.Wrap(err, "setting up logger") + } + conf, err := json.MarshalIndent(m.Config, "", "\t") + if err != nil { + return errors.Wrap(err, "marshalling config") + } + m.logger.Printf("Config: %s", conf) + + // validateAddrs sets the appropriate values for Bind and Advertise + // based on the inputs. It is not responsible for applying defaults, although + // it does provide a non-zero port (10101) in the case where no port is specified. + // The alternative would be to use port 0, which would choose a random port, but + // currently that's not what we want. + if err := m.Config.validateAddrs(context.Background()); err != nil { + return errors.Wrap(err, "validating addresses") + } + + uri, err := featurebase.AddressWithDefaults(m.Config.Bind) + if err != nil { + return errors.Wrap(err, "processing bind address") + } + + m.ln, err = getListener(*uri, m.tlsConfig) + if err != nil { + return errors.Wrap(err, "getting listener") + } + + // If port is 0, get auto-allocated port from listener + if uri.Port == 0 { + uri.SetPort(uint16(m.ln.Addr().(*net.TCPAddr).Port)) + } + + // Save listenURI for later reference. + m.listenURI = uri + + //m.Config.FeatureBase.Config.Listener = ln + + // Get advertise address as uri. + m.advertiseURI, err = featurebase.AddressWithDefaults(m.Config.Advertise) + if err != nil { + return errors.Wrap(err, "processing advertise address") + } + if m.advertiseURI.Port == 0 { + m.advertiseURI.SetPort(uri.Port) + } + + handlerOpts := []daxhttp.HandlerOption{ + daxhttp.OptHandlerBind(m.Config.Bind), + daxhttp.OptHandlerListener(m.ln, m.advertiseURI.String()), + daxhttp.OptHandlerLogger(m.logger), + } + + // Set up WriteLogger. + var wlSvc *writelogger.WriteLogger + if m.Config.WriteLogger.Run { + wlSvc = writelogger.New(writelogger.Config{ + DataDir: m.Config.WriteLogger.Config.DataDir, + Logger: m.logger, + }) + handlerOpts = append(handlerOpts, daxhttp.OptHandlerWriteLogger(wlSvc)) + } + + // Set up Snapshotter. + var ssSvc *snapshotter.Snapshotter + if m.Config.Snapshotter.Run { + ssSvc = snapshotter.New(snapshotter.Config{ + DataDir: m.Config.Snapshotter.Config.DataDir, + Logger: m.logger, + }) + handlerOpts = append(handlerOpts, daxhttp.OptHandlerSnapshotter(ssSvc)) + } + + // Set up MDS. + var mdsSvc *mds.MDS + + // alphaDirector is used in the case where both the `mds` and `computer` + // services are running in the same process. It maintains the mapping + // between computer address and its API. + alphaDirector := controlleralpha.NewDirector() + + alphaRouter := queryeralpha.NewRouter() + + if m.Config.MDS.Run { + mdsSvcCfg := mds.Config{ + RegistrationBatchTimeout: m.Config.MDS.Config.RegistrationBatchTimeout, + StorageMethod: m.Config.StorageMethod, + StorageDSN: m.Config.StorageDSN, + Logger: m.logger, + } + + // If the computer service is being run locally (in process) with MDS, + // then we want to use an implementation of the controller.Director + // interface which calls the interface methods *directly* on the compute + // node service (as opposed to going over http). + if m.Config.Computer.Run { + mdsSvcCfg.Director = alphaDirector + } else { + mdsSvcCfg.Director = controllerhttp.NewDirector( + controllerhttp.DirectorConfig{ + DirectivePath: dax.ServicePrefixComputer + "/directive", + SnapshotRequestPath: dax.ServicePrefixComputer + "/snapshot", + Logger: m.logger, + }) + } + + mdsSvc = mds.New(mdsSvcCfg) + handlerOpts = append(handlerOpts, daxhttp.OptHandlerMDS(mdsSvc)) + + // Start mds services. + if err := mdsSvc.Run(); err != nil { + return errors.Wrap(err, "running mds") + } + } + + // Set up Queryer. + if m.Config.Queryer.Run { + qryrSvcCfg := queryer.Config{ + Logger: m.logger, + } + + var qryrSvcMDS queryer.MDS + var mdsRunning bool + + // This intentionally gives precedence to an MDSAddress over an MDS + // sub-service running in the same process. + if m.Config.Queryer.Config.MDSAddress != "" { + qryrSvcMDS = mdsclient.New(dax.Address(m.Config.Queryer.Config.MDSAddress)) + } else if m.Config.MDS.Run { + qryrSvcMDS = mdsSvc + mdsRunning = true + } else { + return errors.Errorf("queryer can't run without MDS") + } + qryrSvcCfg.MDS = qryrSvcMDS + + // If the computer service is being run locally (in process) with MDS, + // then we want to use an importer which bypasses http requests and + // instead calls the respective services directly. + if mdsRunning && m.Config.Computer.Run { + qryrSvcCfg.Router = alphaRouter + } + + qryrSvc := queryer.New(qryrSvcCfg) + handlerOpts = append(handlerOpts, daxhttp.OptHandlerQueryer(qryrSvc)) + } + + // Set up Computer. + if m.Config.Computer.Run { + // Set the FeatureBase.Config values based on the top-level Config + // values. + m.Config.Computer.Config.Listener = m.ln + m.Config.Computer.Config.Bind = uri.HostPort() + m.Config.Computer.Config.Advertise = m.advertiseURI.HostPort() + + var mdsImpl featurebase.MDS + if m.Config.Computer.Config.MDSAddress != "" { + mdsImpl = mdsclient.New(dax.Address(m.Config.Computer.Config.MDSAddress)) + } else if mdsSvc != nil { + mdsImpl = mdsSvc + } else { + return errors.Errorf("computer requires MDS") + } + + var writeLoggerImpl featurebase.WriteLogger + if m.Config.Computer.Config.WriteLogger != "" { + writeLoggerImpl = writeloggerclient.New(dax.Address(m.Config.Computer.Config.WriteLogger)) + } else if wlSvc != nil { + writeLoggerImpl = wlSvc + } else { + m.logger.Warnf("No writelogger configured, dynamic scaling will not function properly.") + } + + var snapshotterImpl featurebase.Snapshotter + if m.Config.Computer.Config.Snapshotter != "" { + snapshotterImpl = snapshotterclient.New(dax.Address(m.Config.Computer.Config.Snapshotter)) + } else if ssSvc != nil { + snapshotterImpl = ssSvc + } else { + m.logger.Warnf("No snapshotter configured.") + } + + fbcmd := featurebaseserver.NewCommand(m.CmdIO.Stdin, m.CmdIO.Stdout, m.CmdIO.Stderr, + featurebaseserver.OptCommandSetConfig(&m.Config.Computer.Config), + featurebaseserver.OptCommandServerOptions( + featurebase.OptServerIsComputeNode(true), + featurebase.OptServerLogger(m.logger), + ), + featurebaseserver.OptCommandInjections(featurebaseserver.Injections{ + MDS: mdsImpl, + WriteLogger: writeLoggerImpl, + Snapshotter: snapshotterImpl, + IsComputeNode: true, + }), + ) + + // Register the API with the local Director. + if err := alphaDirector.AddCmd(dax.Address(m.advertiseURI.HostPort()), fbcmd); err != nil { + return errors.Wrap(err, "adding cmd to director") + } + + // Register the API with the local Router. + if err := alphaRouter.AddCmd(dax.Address(m.advertiseURI.HostPort()), fbcmd); err != nil { + return errors.Wrap(err, "adding cmd to router") + } + + if err := fbcmd.StartNoServe(); err != nil { + return errors.Wrap(err, "start featurebase command") + } + + // Add the cmd.Register function to the list of functions to call after + // setup. + m.registerFns = append(m.registerFns, fbcmd.Register) + m.checkInFn = fbcmd.CheckIn + + handlerOpts = append(handlerOpts, daxhttp.OptHandlerComputer(fbcmd.HTTPHandler())) + } + + // Set up Handler based on which services are running in process. + m.Handler, err = daxhttp.NewHandler(handlerOpts...) + if err != nil { + return errors.Wrap(err, "new handler") + } + + return nil +} + +// setupLogger sets up the logger based on the configuration. +func (m *Command) setupLogger() error { + var f *logger.FileWriter + var err error + if m.Config.LogPath == "" { + m.logOutput = m.Stderr + } else { + f, err = logger.NewFileWriter(m.Config.LogPath) + if err != nil { + return errors.Wrap(err, "opening file") + } + m.logOutput = f + } + if m.Config.Verbose { + m.logger = logger.NewVerboseLogger(m.logOutput) + } else { + m.logger = logger.NewStandardLogger(m.logOutput) + } + if m.Config.LogPath != "" { + sighup := make(chan os.Signal, 1) + signal.Notify(sighup, syscall.SIGHUP) + go func() { + for { + // duplicate stderr onto log file + err := m.dup(int(f.Fd()), int(os.Stderr.Fd())) + if err != nil { + m.logger.Errorf("syscall dup: %s\n", err.Error()) + } + + // reopen log file on SIGHUP + <-sighup + err = f.Reopen() + if err != nil { + m.logger.Infof("reopen: %s\n", err.Error()) + } + } + }() + } + return nil +} + +// getListener gets a net.Listener based on the config. +func getListener(uri fbnet.URI, tlsconf *tls.Config) (ln net.Listener, err error) { + // If bind URI has the https scheme, enable TLS + if uri.Scheme == "https" && tlsconf != nil { + ln, err = tls.Listen("tcp", uri.HostPort(), tlsconf) + if err != nil { + return nil, errors.Wrap(err, "tls.Listener") + } + } else if uri.Scheme == "http" { + // Open HTTP listener to determine port (if specified as :0). + ln, err = net.Listen("tcp", uri.HostPort()) + if err != nil { + return nil, errors.Wrap(err, "net.Listen") + } + } else { + return nil, errors.Errorf("unsupported scheme: %s", uri.Scheme) + } + + return ln, nil +} diff --git a/dax/shard.go b/dax/shard.go new file mode 100644 index 000000000..2d0264aa5 --- /dev/null +++ b/dax/shard.go @@ -0,0 +1,64 @@ +package dax + +import "fmt" + +// ShardNum is the numerical (uint64) shard value. +type ShardNum uint64 + +// ShardNums is a slice of ShardNum. +type ShardNums []ShardNum + +func (s ShardNum) String() string { + return fmt.Sprintf("%d", s) +} + +// Shard is a versioned shard. +type Shard struct { + Num ShardNum `json:"num"` + Version int `json:"version"` +} + +// NewShard returns a Shard with the provided num and version. +func NewShard(num ShardNum, version int) Shard { + return Shard{ + Num: num, + Version: version, + } +} + +// String returns the Shard (i.e. its Num and Version) as a string. +func (s Shard) String() string { + return fmt.Sprintf("%d.%d", s.Num, s.Version) +} + +// Shards is a sortable slice of Shard. +type Shards []Shard + +func (s Shards) Len() int { return len(s) } +func (s Shards) Less(i, j int) bool { return s[i].Num < s[j].Num } +func (s Shards) Swap(i, j int) { s[i], s[j] = s[j], s[i] } + +// NewShards returns the provided list of shard nums as a list of Shard with an +// invalid version (-1). This is to use for cases where the request should not +// be aware of shard versioning. +func NewShards(shardNums ...ShardNum) Shards { + svs := make(Shards, len(shardNums)) + + for i := range shardNums { + svs[i] = Shard{ + Num: shardNums[i], + Version: -1, + } + } + + return svs +} + +// Nums returns a slice of all the shard numbers in Shards. +func (s Shards) Nums() []ShardNum { + ss := make([]ShardNum, len(s)) + for i := range s { + ss[i] = s[i].Num + } + return ss +} diff --git a/dax/snapshot.go b/dax/snapshot.go new file mode 100644 index 000000000..dd2d06a9c --- /dev/null +++ b/dax/snapshot.go @@ -0,0 +1,34 @@ +package dax + +type SnapshotShardDataRequest struct { + Address Address `json:"address"` + + TableKey TableKey `json:"table-key"` + ShardNum ShardNum `json:"shard"` + FromVersion int `json:"from-version"` + ToVersion int `json:"to-version"` + + Directive Directive `json:"directive"` +} + +type SnapshotTableKeysRequest struct { + Address Address `json:"address"` + + TableKey TableKey `json:"table-key"` + PartitionNum PartitionNum `json:"partition"` + FromVersion int `json:"from-version"` + ToVersion int `json:"to-version"` + + Directive Directive `json:"directive"` +} + +type SnapshotFieldKeysRequest struct { + Address Address `json:"address"` + + TableKey TableKey `json:"table-key"` + Field FieldName `json:"field"` + FromVersion int `json:"from-version"` + ToVersion int `json:"to-version"` + + Directive Directive `json:"directive"` +} diff --git a/dax/snapshotter/api/openapi.yaml b/dax/snapshotter/api/openapi.yaml new file mode 100644 index 000000000..bfb5072ed --- /dev/null +++ b/dax/snapshotter/api/openapi.yaml @@ -0,0 +1,101 @@ +openapi: 3.0.3 + +info: + title: Snapshotter + description: The alpha implementation of the Snapshotter interface. + version: 0.0.0 + +paths: + /snapshotter/health: + get: + summary: Health check endpoint. + description: Provides an endpoint to check the overall health of the Snapshotter service. + operationId: GetHealth + responses: + 200: + description: Service is healthy. + + + /snapshotter/write-snapshot: + post: + summary: Write snapshot. + description: Write snapshot based on bucket/key. + operationId: PostWriteSnapshot + parameters: + - name: bucket + in: query + description: bucket containing snapshot key + required: true + schema: + type: string + - name: key + in: query + description: key identifying snapshot + required: true + schema: + type: string + - name: version + in: query + description: bucket/key version + required: true + schema: + type: integer + format: int64 + requestBody: + content: + text/plain: + schema: + type: string + format: byte + responses: + 200: + $ref: '#/components/responses/WriteSnapshotResponse' + + /snapshotter/read-snapshot: + get: + summary: Read snapshot. + description: Read snapshot based on bucket/key. + operationId: GetReadSnapshot + parameters: + - name: bucket + in: query + description: bucket containing snapshot key + required: true + schema: + type: string + - name: key + in: query + description: key identifying snapshot + required: true + schema: + type: string + - name: version + in: query + description: bucket/key version + required: true + schema: + type: integer + format: int64 + requestBody: + content: + text/plain: + schema: + type: string + format: byte + responses: + 200: + description: Bytes making up the contents of the snapshot. + content: + text/plain: + schema: + type: string + format: byte + +components: + responses: + WriteSnapshotResponse: + description: Placeholder response. + content: + application/json: + schema: + type: object \ No newline at end of file diff --git a/dax/snapshotter/client/client.go b/dax/snapshotter/client/client.go new file mode 100644 index 000000000..852a479fd --- /dev/null +++ b/dax/snapshotter/client/client.go @@ -0,0 +1,109 @@ +// Package client contains an http implementation of the WriteLogger client. +package client + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" + + "github.com/molecula/featurebase/v3/dax" + snapshotterhttp "github.com/molecula/featurebase/v3/dax/snapshotter/http" + "github.com/molecula/featurebase/v3/errors" +) + +const defaultScheme = "http" + +// Snapshotter is a client for the Snapshotter API methods. +type Snapshotter struct { + address dax.Address +} + +func New(address dax.Address) *Snapshotter { + return &Snapshotter{ + address: address, + } +} + +func (s *Snapshotter) Write(bucket string, key string, version int, rc io.ReadCloser) error { + url := fmt.Sprintf("%s/snapshotter/write-snapshot?bucket=%s&key=%s&version=%d", + s.address.WithScheme(defaultScheme), + url.QueryEscape(bucket), + url.QueryEscape(key), + version, + ) + + // Post the request. + resp, err := http.Post(url, "", rc) + if err != nil { + return errors.Wrap(err, "posting write-snapshot") + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(resp.Body) + return errors.Errorf("status code: %d: %s", resp.StatusCode, b) + } + + var wsr snapshotterhttp.WriteSnapshotResponse + if err := json.NewDecoder(resp.Body).Decode(&wsr); err != nil { + return errors.Wrap(err, "reading response body") + } + + return nil +} + +// WriteTo is exactly the same as Write, except that it takes an io.WriteTo +// instead of an io.ReadCloser. This needs to be cleaned up so that we're only +// using one or the other. +func (s *Snapshotter) WriteTo(bucket string, key string, version int, wrTo io.WriterTo) error { + url := fmt.Sprintf("%s/snapshotter/write-snapshot?bucket=%s&key=%s&version=%d", + s.address.WithScheme(defaultScheme), + url.QueryEscape(bucket), + url.QueryEscape(key), + version, + ) + + buf := &bytes.Buffer{} + if _, err := wrTo.WriteTo(buf); err != nil { + return errors.Wrap(err, "writing to buffer") + } + + // Post the request. + resp, err := http.Post(url, "", buf) + if err != nil { + return errors.Wrap(err, "posting write-snapshot") + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(resp.Body) + return errors.Errorf("status code: %d: %s", resp.StatusCode, b) + } + + var wsr snapshotterhttp.WriteSnapshotResponse + if err := json.NewDecoder(resp.Body).Decode(&wsr); err != nil { + return errors.Wrap(err, "reading response body") + } + + return nil +} + +func (s *Snapshotter) Read(bucket string, key string, version int) (io.ReadCloser, error) { + url := fmt.Sprintf("%s/snapshotter/read-snapshot?bucket=%s&key=%s&version=%d", + s.address.WithScheme(defaultScheme), + url.QueryEscape(bucket), + url.QueryEscape(key), + version, + ) + + // Get the request. + resp, err := http.Get(url) + if err != nil { + return nil, errors.Wrap(err, "getting read-snapshot") + } + + return resp.Body, nil +} diff --git a/dax/snapshotter/config.go b/dax/snapshotter/config.go new file mode 100644 index 000000000..7679f3da9 --- /dev/null +++ b/dax/snapshotter/config.go @@ -0,0 +1,8 @@ +package snapshotter + +import "github.com/molecula/featurebase/v3/logger" + +type Config struct { + DataDir string `toml:"data-dir"` + Logger logger.Logger `toml:"-"` +} diff --git a/dax/snapshotter/http/handler.go b/dax/snapshotter/http/handler.go new file mode 100644 index 000000000..82489a765 --- /dev/null +++ b/dax/snapshotter/http/handler.go @@ -0,0 +1,120 @@ +package http + +import ( + "encoding/json" + "io" + "net/http" + "strconv" + + "github.com/gorilla/mux" + "github.com/molecula/featurebase/v3/dax/snapshotter" + "github.com/molecula/featurebase/v3/rbf" +) + +func Handler(s *snapshotter.Snapshotter) http.Handler { + svr := &server{ + snapshotter: s, + } + + router := mux.NewRouter() + router.HandleFunc("/health", svr.getHealth).Methods("GET").Name("GetHealth") + router.HandleFunc("/write-snapshot", svr.postWriteSnapshot).Methods("POST").Name("PostWriteSnapshot") + router.HandleFunc("/read-snapshot", svr.getReadSnapshot).Methods("GET").Name("GetReadSnapshot") + return router +} + +type server struct { + snapshotter *snapshotter.Snapshotter +} + +// GET /health +func (s *server) getHealth(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) +} + +// POST /write-snapshot +func (s *server) postWriteSnapshot(w http.ResponseWriter, r *http.Request) { + bucket := r.URL.Query().Get("bucket") + if bucket == "" { + http.Error(w, "bucket required", http.StatusBadRequest) + return + } + + key := r.URL.Query().Get("key") + if key == "" { + http.Error(w, "key required", http.StatusBadRequest) + return + } + + versionArg := r.URL.Query().Get("version") + versionInt64, err := strconv.ParseInt(versionArg, 10, 64) + if err != nil { + http.Error(w, "bad shard", http.StatusBadRequest) + return + } + version := int(versionInt64) + + body := r.Body + defer body.Close() + + if err := s.snapshotter.Write(bucket, key, version, body); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + resp := &WriteSnapshotResponse{} + + if err := json.NewEncoder(w).Encode(resp); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } +} + +type WriteSnapshotResponse struct{} + +// GET /read-snapshot +func (s *server) getReadSnapshot(w http.ResponseWriter, r *http.Request) { + bucket := r.URL.Query().Get("bucket") + if bucket == "" { + http.Error(w, "bucket required", http.StatusBadRequest) + return + } + + key := r.URL.Query().Get("key") + if key == "" { + http.Error(w, "key required", http.StatusBadRequest) + return + } + + versionArg := r.URL.Query().Get("version") + versionInt64, err := strconv.ParseInt(versionArg, 10, 64) + if err != nil { + http.Error(w, "bad shard", http.StatusBadRequest) + return + } + version := int(versionInt64) + + rc, err := s.snapshotter.Read(bucket, key, version) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + defer rc.Close() + + // TODO: is rbf.PageSize a problem here for non-RBF snapshots (i.e. keys)? + // Copy data to response body. + if _, err := io.CopyBuffer(&passthroughWriter{w}, rc, make([]byte, rbf.PageSize)); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } +} + +// passthroughWriter is used to remove non-Writer interfaces from an io.Writer. +// For example, a writer that implements io.ReaderFrom can change io.Copy() behavior. +type passthroughWriter struct { + w io.Writer +} + +func (w *passthroughWriter) Write(p []byte) (int, error) { + return w.w.Write(p) +} diff --git a/dax/snapshotter/snapshotter.go b/dax/snapshotter/snapshotter.go new file mode 100644 index 000000000..ad7b5a8ee --- /dev/null +++ b/dax/snapshotter/snapshotter.go @@ -0,0 +1,112 @@ +// Package snapshotter provides the core snapshotter structs. +package snapshotter + +import ( + "bytes" + "fmt" + "io" + "io/fs" + "os" + "path" + "sync" + + "github.com/molecula/featurebase/v3/errors" + "github.com/molecula/featurebase/v3/logger" +) + +type Snapshotter struct { + mu sync.RWMutex + + dataDir string + + logger logger.Logger +} + +func New(cfg Config) *Snapshotter { + return &Snapshotter{ + dataDir: cfg.DataDir, + logger: logger.NopLogger, + } +} + +// SetLogger sets the logger used for logging messages. +func (s *Snapshotter) SetLogger(l logger.Logger) { + s.logger = l +} + +func (s *Snapshotter) Write(bucket string, key string, version int, rc io.ReadCloser) error { + fKey := fullKey(bucket, key, version) + snapshotFile, err := s.snapshotFileByKey(fKey) + if err != nil { + return errors.Wrapf(err, "shapshotting file by key: %s", fKey) + } + defer snapshotFile.Close() + + defer rc.Close() + if _, err := snapshotFile.ReadFrom(rc); err != nil { + return errors.Wrap(err, "reading from shapshot file") + } + + return snapshotFile.Sync() +} + +func (s *Snapshotter) Read(bucket string, key string, version int) (io.ReadCloser, error) { + _, filePath := s.paths(fullKey(bucket, key, version)) + f, err := os.Open(filePath) + if err != nil { + if e, ok := err.(*fs.PathError); ok { + return nil, e + } + return nil, errors.Wrapf(err, "reading snapshot file: %s", filePath) + } + + return f, nil +} + +// WriteTo is exactly the same as Write, except that it takes an io.WriteTo +// instead of an io.ReadCloser. This needs to be cleaned up so that we're only +// using one or the other. +func (s *Snapshotter) WriteTo(bucket string, key string, version int, wrTo io.WriterTo) error { + buf := &bytes.Buffer{} + if _, err := wrTo.WriteTo(buf); err != nil { + return errors.Wrap(err, "writing to buffer") + } + return s.Write(bucket, key, version, io.NopCloser(buf)) +} + +// paths takes a key and returns the full file path (including the root data +// directory) as well as the full directory path (i.e. the file path without the +// file portion). +func (s *Snapshotter) paths(key string) (string, string) { + filePath := path.Join(s.dataDir, key) + dirPath, _ := path.Split(filePath) + return dirPath, filePath +} + +// snapshotFileByKey returns a pointer to the file specified by key. If the file +// does not exist, the file is created (along with any directories in which the +// file is nested). +func (s *Snapshotter) snapshotFileByKey(key string) (*os.File, error) { + s.mu.Lock() + defer s.mu.Unlock() + + dirPath, filePath := s.paths(key) + + // make directories + if err := os.MkdirAll(dirPath, 0777); err != nil { + return nil, errors.Wrapf(err, "making directory: %s", dirPath) + } + + // open snapshot file + f, err := os.OpenFile(filePath, os.O_CREATE|os.O_WRONLY, 0644) + if err != nil { + return nil, errors.Wrapf(err, "opening shapshot file: %s", filePath) + } + + return f, nil +} + +// fullKey returns the full file key including the bucket and version. +func fullKey(bucket string, key string, version int) string { + return path.Join(bucket, key, fmt.Sprintf("%d", version)) +} diff --git a/dax/snapshotter/snapshotter_test.go b/dax/snapshotter/snapshotter_test.go new file mode 100644 index 000000000..7aa0f5193 --- /dev/null +++ b/dax/snapshotter/snapshotter_test.go @@ -0,0 +1,70 @@ +package snapshotter_test + +import ( + "fmt" + "os" + "path" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestSnapshotter(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "testWriteLogger-*") + assert.NoError(t, err) + + // Remove the temp directory. + defer func() { + os.RemoveAll(tmpDir) + }() + + // t.Run("Basic", func(t *testing.T) { + // type payload struct { + // Foo string `json:"foo"` + // Bar int `json:"bar"` + // } + + // cfg := core.Config{ + // DataDir: tmpDir, + // } + // wl := core.NewSnapshotter(cfg) + + // table := "tbl" + // partition := 1 + // version := 0 + // key := "keys" + + // msg1 := payload{ + // Foo: "message 1", + // Bar: 88, + // } + + // // Write the message. + // msg, err := json.Marshal(msg1) + // assert.NoError(t, err) + + // err = wl.AppendMessage(bucket(table, partition), key, version, msg) + // assert.NoError(t, err) + + // // Read the message. + // reader, closer, err := wl.LogReader(bucket(table, partition), key, version) + // assert.NoError(t, err) + // defer closer.Close() + + // buf, err := ioutil.ReadAll(reader) + // assert.NoError(t, err) + + // var out payload + + // err = json.Unmarshal(buf, &out) + // assert.NoError(t, err) + + // assert.Equal(t, msg1.Foo, out.Foo) + // assert.Equal(t, msg1.Bar, out.Bar) + // }) +} + +func bucket(table string, partition int) string { + return path.Join(table, fmt.Sprintf("%d", partition)) + +} diff --git a/dax/table.go b/dax/table.go new file mode 100644 index 000000000..234017560 --- /dev/null +++ b/dax/table.go @@ -0,0 +1,581 @@ +package dax + +import ( + "crypto/rand" + "fmt" + "strings" + "time" + + "github.com/molecula/featurebase/v3/errors" + "github.com/molecula/featurebase/v3/pql" +) + +//////////////////////////////////////////////////////////////////////////////// +// +// Table +// +// The types defined below are used to standardize on tables and fields. Prior +// to introducing these types, the only way we could identify a table was by +// name, which wasn't even a defined type. Rather, we passed `string` values +// throughout the code. +// +// OrganizationID - carried over from ControlPlane; currently uuid +// DatabaseID - carried over from ControlPlane; currently uuid +// TableID - internally stored as a uint64; presented as a hex string. +// TableName - human-friendly string table name +// Table - base Table struct; includes a TableID and a TableName +// TableQualifier - combination of OrganizationID and DatabaseID +// QualifiedTable - TableQualifier plus a Table +// QualifiedTableID - TableQualifer plus a TableID +// TableKey - a string representation of OrganizationID, DatabaseID, and +// TableID, which is safe to use as a FeatureBase index name. +// +// Example: +// OrganizationID - "29-ae44-41" +// DatabaseID - "75-d1a2-4f" +// TableID - 123456789 (hex string: "499602d2") +// TableName - foo +// Table - {ID:"499602d2", Name: "foo", Fields: ... } +// TableQualifier - {Org: "29-ae44-41", DB: "75-d1a2-4f"} +// QualifiedTable - {Org: "29-ae44-41", DB: "75-d1a2-4f", Table: *tbl} +// QualifiedTableID - {Org: "29-ae44-41", DB: "75-d1a2-4f", TableID: "499602d2"} +// TableKey - "tbl__29-ae44-41__75-d1a2-4f__499602d2" +// +//////////////////////////////////////////////////////////////////////////////// + +// TableKeyDelimiter is used to delimit the qualifer elements in the TableKey. +// While it might make more sense to use a pipe ("|") here, we instead use a +// double underscore because underscore is one of the few characters allowed by +// the FeatureBase index name restrictions, and we double it in a lame attempt +// to distinquish it from FeatureBase index names which contain a single +// underscore. +const TableKeyDelimiter = "__" + +// PrefixTable is used as a prefix to TableKey strings because FeatureBase +// indexes must start with an alpha (a-z) character. Because the string +// representation of a uuid (i.e. the OrganizationID value) can start with a +// numeric value, we can't have OrganizationId (or any of the other ID values +// which make up the TableKey) be at the beginning of the TableKey. +const PrefixTable = "tbl" + +// Field types. +const ( + FieldTypeBool = "bool" // + FieldTypeDecimal = "decimal" // + FieldTypeID = "id" // non-keyed mutex + FieldTypeIDSet = "idset" // non-keyed set + FieldTypeInt = "int" // + FieldTypeString = "string" // keyed mutex + FieldTypeStringSet = "stringset" // keyed set + FieldTypeTimestamp = "timestamp" // + + DefaultPartitionN = 256 + + PrimaryKeyFieldName = FieldName("_id") +) + +// Schema contains a list of Tables. +type Schema struct { + Tables []*Table +} + +// Table returns the table with the provided name. If a table with that name +// does not exist, the returned boolean will be false. +func (s *Schema) Table(name TableName) (*Table, bool) { + for _, tbl := range s.Tables { + if tbl.Name == name { + return tbl, true + } + } + return nil, false +} + +// OrganizationID is the unique organization identifier, currently generated by +// the Control Plane in a FeatureBase cloud implementation. In that +// implementation, its value is a uuid as a string, but there's nothing +// enforcing that; the value could be any string. +type OrganizationID string + +// DatabaseID is the unique database identifier, currently generated by the +// Control Plane in a FeatureBase cloud implementation. In that implementation, +// its value is a uuid as a string, but there's nothing enforcing that; the +// value could be any string. +type DatabaseID string + +// TableKey is a globally unique identifier for a table; it is effectively the +// compound key: (org, database, table). This is (hopefully) the value that will +// be used when interfacing with services which are unaware of table qualifiers. +// For example, the FeatureBase server has no notion of organization or +// database; its top level type is index/indexName/table. So in this case, until +// and unless we introduce table qualifiers into FeatureBase, we will use +// TableKey as the value for index.Name. +type TableKey string + +// QualifiedTableID returns the QualifiedTableID based on the key. If TableKey +// can't be parsed into a valid (i.e. complete) QualifiedTableID, then blank +// values are used where necessary. +func (tk TableKey) QualifiedTableID() QualifiedTableID { + qtid, err := QualifiedTableIDFromKey(string(tk)) + if err != nil { + return NewQualifiedTableID( + NewTableQualifier("", ""), + TableID(tk), + ) + } + return qtid +} + +// TableKeys is a sortable slice of TableKey. +type TableKeys []TableKey + +func (s TableKeys) Len() int { return len(s) } +func (s TableKeys) Less(i, j int) bool { return s[i] < s[j] } +func (s TableKeys) Swap(i, j int) { s[i], s[j] = s[j], s[i] } + +// TableID is a table identifier. It is unique within the scope of a +// TableQualifier. Coupled with a TableQualifier, it makes up a +// QualifiedTableID and, when encoded as a string, a TableKey. +type TableID string + +// TableIDs is a sortable slice of TableID. +type TableIDs []TableID + +func (s TableIDs) Len() int { return len(s) } +func (s TableIDs) Less(i, j int) bool { return s[i] < s[j] } +func (s TableIDs) Swap(i, j int) { s[i], s[j] = s[j], s[i] } + +// TableName is a human-friendly string. While it is not used as a primary key, +// uniqueness is generally enforced within the scope of a TableQualifier. +type TableName string + +// TableNames is a sortable slice of TableName. +type TableNames []TableName + +func (s TableNames) Len() int { return len(s) } +func (s TableNames) Less(i, j int) bool { return s[i] < s[j] } +func (s TableNames) Swap(i, j int) { s[i], s[j] = s[j], s[i] } + +// Table represents a table and its configuration. +type Table struct { + ID TableID `json:"id,omitempty"` + Name TableName `json:"name,omitempty"` + Fields []*Field `json:"fields"` + PartitionN int `json:"partitionN"` +} + +// CreateID generates a unique identifier for Table. If Table has already been +// assigned an ID, then an error is returned. +func (t *Table) CreateID() (TableID, error) { + if t.ID != "" { + return "", errors.Errorf("CreateID called on table %+v that already has ID", t) + } + + // stub is prepended to the Table.ID as a way to make IDs somewhat + // human-readable for debugging purposes. If the table name is changed after + // its ID has been created, this could be confusing (because the stub + // portion of the ID will still resemble the initial table name). + // + // In order to avoid creating an ID with a double underscore, we remove all + // underscores from the original table name (because that's what we use in + // TableKey as a delimiter). + stub := strings.ReplaceAll(string(t.Name), "_", "") + if len(stub) > 10 { + stub = stub[:10] + } + + rn := make([]byte, 8) + if _, err := rand.Read(rn); err != nil { + return "", errors.Wrap(err, "getting random data") + } + t.ID = TableID(fmt.Sprintf("%s_%x", stub, rn)) + + return t.ID, nil +} + +// NewTable returns a new instance of table with a pseudo-random ID which is +// assumed to be unique within the scope of a TableQualifer. +func NewTable(name TableName) *Table { + return &Table{ + Name: name, + Fields: make([]*Field, 0), + } +} + +// StringKeys returns true if the table's primary key is either a string or a +// concatenation of fields. +func (t *Table) StringKeys() bool { + for _, fld := range t.Fields { + if fld.IsPrimaryKey() { + if fld.Type == FieldTypeString { + return true + } + break + } + } + return false +} + +// HasValidPrimaryKey returns false if the table does not contain a primary key +// field (which is required), or if the primary key field is not a valid type. +func (t *Table) HasValidPrimaryKey() bool { + for _, fld := range t.Fields { + if !fld.IsPrimaryKey() { + continue + } + + if fld.Type == FieldTypeID || fld.Type == FieldTypeString { + return true + } + } + return false +} + +// FieldNames returns the list of field names associated with the table. +func (t *Table) FieldNames() []FieldName { + var ret []FieldName + for _, f := range t.Fields { + ret = append(ret, f.Name) + } + return ret +} + +// Field returns the field with the provided name. If a field with that name +// does not exist, the returned boolean will be false. +func (t *Table) Field(name FieldName) (*Field, bool) { + for _, fld := range t.Fields { + if fld.Name == name { + return fld, true + } + } + return nil, false +} + +// RemoveField removes the given field by name. It returns true if the field was +// removed. +func (t *Table) RemoveField(name FieldName) bool { + for i, fld := range t.Fields { + if fld.Name == name { + t.Fields = append(t.Fields[:i], t.Fields[i+1:]...) + return true + } + } + return false +} + +// CreateSQL returns the SQL CREATE TABLE string necessary to create the table. +func (t *Table) CreateSQL() string { + sql := fmt.Sprintf("CREATE TABLE %s (", t.Name) + + cols := []string{} + for _, fld := range t.Fields { + cols = append(cols, fld.CreateSQL()) + } + sql += strings.Join(cols, ", ") + + sql += fmt.Sprintf(") KEYPARTITIONS %d", t.PartitionN) + + return sql +} + +// Tables is a sortable slice of Table. +type Tables []*Table + +func (o Tables) Len() int { return len(o) } +func (o Tables) Less(i, j int) bool { return o[i].Name < o[j].Name } +func (o Tables) Swap(i, j int) { o[i], o[j] = o[j], o[i] } + +// TableQualifierKey is the unique TableQualifer values encoded as a string. The +// current encoding is delimited as `prefix|OrganizationID|DatabaseID` (where +// the pipe may be some other delimiter) by the TableQualifier.Key() method. +type TableQualifierKey string + +// Qualifier returns the Qualifier based on the values encoded into the +// TableQualifierKey string. +func (tqk TableQualifierKey) Qualifier() TableQualifier { + parts := strings.Split(string(tqk), TableKeyDelimiter) + + if len(parts) < 3 { + return NewTableQualifier("", "") + } + + return NewTableQualifier( + OrganizationID(parts[1]), + DatabaseID(parts[2]), + ) +} + +// OrganizationID returns the OrganizationID value that has been encoded into +// the TableQualifierKey string. +func (tqk TableQualifierKey) OrganizationID() OrganizationID { + parts := strings.Split(string(tqk), TableKeyDelimiter) + + if len(parts) < 2 { + return "" + } + + return OrganizationID(parts[1]) +} + +// DatabaseID returns the DatabaseID value that has been encoded into the +// TableQualifierKey string. +func (tqk TableQualifierKey) DatabaseID() DatabaseID { + parts := strings.Split(string(tqk), TableKeyDelimiter) + + if len(parts) < 3 { + return "" + } + + return DatabaseID(parts[2]) +} + +// TableQualifier contains all the elements required to fully qualify a table. +type TableQualifier struct { + OrganizationID OrganizationID `json:"org-id"` + DatabaseID DatabaseID `json:"db-id"` +} + +// NewTableQualifier is a helper function used to create a TableQualifer from +// the provided arguments. +func NewTableQualifier(orgID OrganizationID, dbID DatabaseID) TableQualifier { + return TableQualifier{ + OrganizationID: orgID, + DatabaseID: dbID, + } +} + +// String returns a human-friendly version of the TableQualifier. It is only +// used for display purposes; it is not used as any kind of key. For that, see +// the TableQualifier.Key() method and the TableQualifierKey type. +func (tq TableQualifier) String() string { + return fmt.Sprintf("[%s:%s]", tq.OrganizationID, tq.DatabaseID) +} + +// Key returns the string-encoded (delimited by TableKeyDelimiter) +// TableQualifierKey. +func (tq TableQualifier) Key() TableQualifierKey { + return TableQualifierKey(fmt.Sprintf("%s%s%s%s%s", + PrefixTable, + TableKeyDelimiter, + tq.OrganizationID, + TableKeyDelimiter, + tq.DatabaseID, + )) +} + +//////////////////////////////////////////////// + +// QualifiedTableID is a globally unique table identifier. It is a +// sub-set of a QualifiedTable (i.e. it's just the identification +// portion). Most things will take a Name or an ID and do the right +// thing™. +type QualifiedTableID struct { + TableQualifier + ID TableID `json:"id"` + Name TableName `json:"name"` +} + +// NewQualifiedTableID is a helper function used to create a QualifiedTableID +// from the provided arguments. +func NewQualifiedTableID(q TableQualifier, id TableID) QualifiedTableID { + return QualifiedTableID{ + TableQualifier: q, + ID: id, + } +} + +// QualifiedTableIDFromKey decodes a string key into a QualifiedTableID. The key +// is assumed to have been encoded using the QualifiedTableID.Key() method. +func QualifiedTableIDFromKey(key string) (QualifiedTableID, error) { + parts := strings.Split(key, TableKeyDelimiter) + switch len(parts) { + case 4: + // prefix|orgID|dbID|tblID + return NewQualifiedTableID( + NewTableQualifier( + OrganizationID(parts[1]), + DatabaseID(parts[2]), + ), + TableID(parts[3]), + ), nil + default: + return QualifiedTableID{}, errors.Errorf("invalid key: %s", key) + } +} + +// String returns a human-friendly version of the TableQualifierID. It is only +// used for display purposes; it is not used as any kind of key. For that, see +// the TableQualifierID.Key() method. +func (qtid QualifiedTableID) String() string { + if qtid.ID == "" { + return fmt.Sprintf("%s%s", qtid.TableQualifier, qtid.Name) + } + return fmt.Sprintf("%s%s", qtid.TableQualifier, qtid.ID) +} + +// Key returns the string-encoded (delimited by TableKeyDelimiter) globally +// unique TableKey. The key has a prefix because FeatureBase index name +// restrictions require the name to start with a non-numeric value, and since a +// uuid can contain a number as its first character, we have to prefix it with +// something. +func (qtid QualifiedTableID) Key() TableKey { + if qtid.ID == "" { + panic("QualifiedTableID.Key called without an ID set") + } + return TableKey(fmt.Sprintf("%s%s%s", + qtid.TableQualifier.Key(), + TableKeyDelimiter, + qtid.ID)) +} + +// Equals returns true if `other` is the same as qtid. Note: the `Name` value is +// ignored in this comparison; only `TableQaulifer` and `ID` are considered. +func (qtid QualifiedTableID) Equals(other QualifiedTableID) bool { + if qtid.TableQualifier == other.TableQualifier && qtid.ID == other.ID { + return true + } + return false +} + +//////////////////////////////////////////////// + +// QualifiedTable wraps Table and includes a TableQualifier. +type QualifiedTable struct { + Table + TableQualifier +} + +// NewQualifiedTable returns the tbl as a QualifiedTable with the provided +// TableQualifier. +func NewQualifiedTable(qual TableQualifier, tbl *Table) *QualifiedTable { + return &QualifiedTable{ + Table: *tbl, + TableQualifier: qual, + } +} + +// Key returns the string-encoded (delimited by TableKeyDelimiter) globally +// unique TableKey. +func (qt QualifiedTable) Key() TableKey { + return qt.QualifiedID().Key() +} + +// String returns a human-friendly version of the QualifiedTable. It is only +// used for display purposes; it is not used as any kind of key. +func (qt QualifiedTable) String() string { + return fmt.Sprintf("%s (%s)", qt.QualifiedID(), qt.Name) +} + +// Qualifier returns the TableQualifer portion of the QualifiedTable. +func (qt *QualifiedTable) Qualifier() TableQualifier { + return qt.TableQualifier +} + +// QualifiedID returns the QualifiedTableID for the table. +func (qt *QualifiedTable) QualifiedID() QualifiedTableID { + return QualifiedTableID{ + TableQualifier: qt.TableQualifier, + ID: qt.ID, + Name: qt.Name, + } +} + +// QualifiedTables is a sortable slice of QualifiedTable. +type QualifiedTables []*QualifiedTable + +func (o QualifiedTables) Len() int { return len(o) } +func (o QualifiedTables) Less(i, j int) bool { return o[i].ID < o[j].ID } +func (o QualifiedTables) Swap(i, j int) { o[i], o[j] = o[j], o[i] } + +// FieldName is a typed string used for field names. +type FieldName string + +// FieldType is a typed string used for field types. +type FieldType string + +// Field represents a field and its configuration. +type Field struct { + Name FieldName `json:"name"` + Type FieldType `json:"type"` + Options FieldOptions `json:"options"` +} + +// String returns the field name as a string. +func (f *Field) String() string { + return string(f.Name) +} + +// StringKeys returns true if the field uses string keys. +func (f *Field) StringKeys() bool { + switch f.Type { + case FieldTypeString, FieldTypeStringSet: + return true + } + return false +} + +// IsPrimaryKey returns true if the field is the primary key field (of either +// type ID or STRING). +func (f *Field) IsPrimaryKey() bool { + return f.Name == PrimaryKeyFieldName +} + +// CreateSQL returns the SQL representation of the field to be used in a CREATE +// TABLE statement. +func (f *Field) CreateSQL() string { + sql := fmt.Sprintf("%s %s", f.Name, f.Type) + + // Apply constraints to all non-primarykey fields. + if !f.IsPrimaryKey() { + sql += f.constraints() + } + + return sql +} + +func (f *Field) constraints() string { + sql := "" + + // Apply constraints. + switch f.Type { + case FieldTypeInt: + sql += fmt.Sprintf(" MIN %d MAX %d", f.Options.Min.ToInt64(0), f.Options.Max.ToInt64(0)) + case FieldTypeID, FieldTypeString: + if f.Options.CacheType != "" { + sql += fmt.Sprintf(" CACHETYPE %s SIZE %d", f.Options.CacheType, f.Options.CacheSize) + } + case FieldTypeIDSet, FieldTypeStringSet: + if f.Options.CacheType != "" { + sql += fmt.Sprintf(" CACHETYPE %s SIZE %d", f.Options.CacheType, f.Options.CacheSize) + } + if f.Options.TimeQuantum != "" { + sql += fmt.Sprintf(" TIMEQUANTUM '%s'", f.Options.TimeQuantum) + if f.Options.TTL > 0 { + sql += fmt.Sprintf(" TTL '%s'", f.Options.TTL) + } + } + case FieldTypeTimestamp: + if f.Options.TimeUnit != "" { + sql += fmt.Sprintf(" TIMEUNIT '%s'", f.Options.TimeUnit) + if !f.Options.Epoch.IsZero() { + sql += fmt.Sprintf(" EPOCH '%s'", f.Options.Epoch.Format(time.RFC3339)) // time.RFC3339 + } + } + } + + return sql +} + +// FieldOptions represents options to set when initializing a field. +type FieldOptions struct { + Min pql.Decimal `json:"min,omitempty"` + Max pql.Decimal `json:"max,omitempty"` + Scale int64 `json:"scale,omitempty"` + NoStandardView bool `json:"no-standard-view,omitempty"` // TODO: we should remove this + CacheType string `json:"cache-type,omitempty"` + CacheSize uint32 `json:"cache-size,omitempty"` + TimeUnit string `json:"time-unit,omitempty"` + Epoch time.Time `json:"epoch,omitempty"` + TimeQuantum TimeQuantum `json:"time-quantum,omitempty"` + TTL time.Duration `json:"ttl,omitempty"` + ForeignIndex string `json:"foreign-index,omitempty"` +} diff --git a/dax/table_test.go b/dax/table_test.go new file mode 100644 index 000000000..4043c0871 --- /dev/null +++ b/dax/table_test.go @@ -0,0 +1,325 @@ +package dax_test + +import ( + "encoding/json" + "fmt" + "sort" + "testing" + "time" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/pql" + "github.com/stretchr/testify/assert" +) + +func TestTable(t *testing.T) { + tableName := dax.TableName("foo") + + t.Run("StringKeys", func(t *testing.T) { + // No PrimaryKeys. + { + tbl := &dax.Table{ + Name: tableName, + } + assert.False(t, tbl.StringKeys()) + } + // One PrimaryKey (string). + { + tbl := &dax.Table{ + Name: tableName, + Fields: []*dax.Field{ + { + Name: dax.PrimaryKeyFieldName, + Type: dax.FieldTypeString, + }, + { + Name: "stringField2", + Type: dax.FieldTypeString, + }, + }, + } + assert.True(t, tbl.StringKeys()) + } + // One PrimaryKey (non-string). + { + tbl := &dax.Table{ + Name: tableName, + Fields: []*dax.Field{ + { + Name: dax.PrimaryKeyFieldName, + Type: dax.FieldTypeID, + }, + }, + } + assert.False(t, tbl.StringKeys()) + } + }) + + t.Run("Tables", func(t *testing.T) { + tblA := &dax.Table{ + Name: "a", + } + tblZ := &dax.Table{ + Name: "z", + } + + tables := []*dax.Table{tblZ, tblA} + + // Ensure tables starts out of order. + assert.Equal(t, tblZ.Name, tables[0].Name) + assert.Equal(t, tblA.Name, tables[1].Name) + + sort.Sort(dax.Tables(tables)) + + // Ensure tables is ordered by name. + assert.Equal(t, tblA.Name, tables[0].Name) + assert.Equal(t, tblZ.Name, tables[1].Name) + + }) + + t.Run("CreateSQL", func(t *testing.T) { + tests := []struct { + tbl dax.Table + expSQL string + }{ + { + tbl: dax.Table{ + Name: "just_a_table", + Fields: []*dax.Field{ + { + Name: "_id", + Type: "id", + }, + }, + }, + expSQL: "CREATE TABLE just_a_table (_id id) KEYPARTITIONS 0", + }, + { + tbl: dax.Table{ + Name: "all_field_types", + Fields: []*dax.Field{ + { + Name: "_id", + Type: "string", + }, + { + Name: "an_id", + Type: "id", + }, + { + Name: "a_string", + Type: "string", + }, + { + Name: "an_id_set", + Type: "idset", + }, + { + Name: "a_string_set", + Type: "stringset", + }, + { + Name: "an_int", + Type: "int", + }, + { + Name: "a_decimal", + Type: "decimal", + }, + { + Name: "a_timestamp", + Type: "timestamp", + }, + }, + }, + expSQL: "CREATE TABLE all_field_types (_id string, an_id id, a_string string, an_id_set idset, a_string_set stringset, an_int int MIN 0 MAX 0, a_decimal decimal, a_timestamp timestamp) KEYPARTITIONS 0", + }, + { + tbl: dax.Table{ + Name: "all_field_types_with_options", + Fields: []*dax.Field{ + { + Name: "_id", + Type: "string", + }, + { + Name: "an_id", + Type: "id", + Options: dax.FieldOptions{ + CacheType: "ranked", + CacheSize: 500, + }, + }, + { + Name: "a_string", + Type: "string", + Options: dax.FieldOptions{ + CacheType: "ranked", + CacheSize: 500, + }, + }, + { + Name: "an_id_set", + Type: "idset", + Options: dax.FieldOptions{ + CacheType: "ranked", + CacheSize: 500, + }, + }, + { + Name: "a_string_set", + Type: "stringset", + Options: dax.FieldOptions{ + CacheType: "ranked", + CacheSize: 500, + }, + }, + { + Name: "an_int", + Type: "int", + Options: dax.FieldOptions{ + Min: pql.NewDecimal(-100, 0), + Max: pql.NewDecimal(200, 0), + }, + }, + { + Name: "a_decimal", + Type: "decimal", + Options: dax.FieldOptions{}, + }, + { + Name: "a_timestamp", + Type: "timestamp", + Options: dax.FieldOptions{ + TimeUnit: "s", + Epoch: time.Date(2009, 11, 10, 23, 34, 56, 0, time.UTC), + }, + }, + }, + }, + expSQL: "CREATE TABLE all_field_types_with_options (_id string, an_id id CACHETYPE ranked SIZE 500, a_string string CACHETYPE ranked SIZE 500, an_id_set idset CACHETYPE ranked SIZE 500, a_string_set stringset CACHETYPE ranked SIZE 500, an_int int MIN -100 MAX 200, a_decimal decimal, a_timestamp timestamp TIMEUNIT 's' EPOCH '2009-11-10T23:34:56Z') KEYPARTITIONS 0", + }, + } + for i, test := range tests { + t.Run(fmt.Sprintf("test-%d", i), func(t *testing.T) { + assert.Equal(t, test.expSQL, test.tbl.CreateSQL()) + }) + } + }) + + t.Run("Table", func(t *testing.T) { + t.Run("RandomID", func(t *testing.T) { + n := dax.NewTable(tableName) + assert.Empty(t, n.ID) + n.CreateID() + assert.NotEmpty(t, n.ID) + assert.Equal(t, tableName, n.Name) + }) + + t.Run("ToJSON", func(t *testing.T) { + t.Run("WithID", func(t *testing.T) { + n := dax.NewTable(tableName) + n.CreateID() + id := n.ID + assert.NotEmpty(t, id) + + b, err := json.Marshal(n) + assert.NoError(t, err) + + exp := fmt.Sprintf("{\"id\":\"%s\",\"name\":\"%s\",\"fields\":[],\"partitionN\":0}", id, tableName) + assert.JSONEq(t, exp, string(b)) + }) + + t.Run("WithoutID", func(t *testing.T) { + n := dax.Table{ + Name: tableName, + } + id := n.ID + assert.Empty(t, id) + + b, err := json.Marshal(n) + assert.NoError(t, err) + + exp := fmt.Sprintf("{\"name\":\"%s\",\"fields\":null,\"partitionN\":0}", tableName) + assert.JSONEq(t, exp, string(b)) + }) + }) + + t.Run("FromJSON", func(t *testing.T) { + tid := dax.TableID("0000000000abc123") // 11256099 + j := fmt.Sprintf("{\"id\":\"%s\",\"name\":\"foo\",\"fields\":[],\"partitionN\":0}", tid) + + tbl := &dax.Table{} + err := json.Unmarshal([]byte(j), tbl) + assert.NoError(t, err) + assert.Equal(t, tid, tbl.ID) + }) + }) + + t.Run("QualifiedTable", func(t *testing.T) { + orgID := dax.OrganizationID("acme") + dbID := dax.DatabaseID("db1") + + t.Run("New", func(t *testing.T) { + tbl := dax.NewTable(tableName) + tbl.CreateID() + qual := dax.TableQualifier{ + OrganizationID: orgID, + DatabaseID: dbID, + } + qtbl := dax.NewQualifiedTable(qual, tbl) + assert.NotEmpty(t, qtbl.ID) + assert.Equal(t, tbl.ID, qtbl.ID) + + tq := qtbl.Qualifier() + assert.Equal(t, qual.OrganizationID, tq.OrganizationID) + assert.Equal(t, qual.DatabaseID, tq.DatabaseID) + + wrappedTable := qtbl.Table + assert.Equal(t, tbl.Name, wrappedTable.Name) + assert.Equal(t, tbl.ID, wrappedTable.ID) + + // Key. + exp := dax.TableKey(fmt.Sprintf("%s%s%s%s%s%s%s", + dax.PrefixTable, + dax.TableKeyDelimiter, + orgID, + dax.TableKeyDelimiter, + dbID, + dax.TableKeyDelimiter, + tbl.ID, + )) + assert.Equal(t, exp, qtbl.Key()) + }) + + t.Run("ToJSON", func(t *testing.T) { + tbl := dax.NewTable(tableName) + tbl.CreateID() + qual := dax.TableQualifier{ + OrganizationID: orgID, + DatabaseID: dbID, + } + qtbl := dax.NewQualifiedTable(qual, tbl) + id := qtbl.ID + + b, err := json.Marshal(qtbl) + assert.NoError(t, err) + + exp := fmt.Sprintf("{\"org-id\":\"%s\",\"db-id\":\"%s\",\"id\":\"%s\",\"name\":\"%s\",\"fields\":[],\"partitionN\":0}", orgID, dbID, id, tableName) + assert.JSONEq(t, exp, string(b)) + }) + + t.Run("FromJSON", func(t *testing.T) { + tid := dax.TableID("0000000000abc123") // 11256099 + j := fmt.Sprintf("{\"org-id\":\"%s\",\"db-id\":\"%s\",\"id\":\"%s\",\"name\":\"foo\",\"fields\":[],\"partitionN\":0}", orgID, dbID, tid) + + qtbl := &dax.QualifiedTable{} + err := json.Unmarshal([]byte(j), qtbl) + assert.NoError(t, err) + + assert.Equal(t, tid, qtbl.ID) + assert.Equal(t, orgID, qtbl.Qualifier().OrganizationID) + assert.Equal(t, dbID, qtbl.Qualifier().DatabaseID) + }) + }) +} diff --git a/dax/test/boltdb/helpers.go b/dax/test/boltdb/helpers.go new file mode 100644 index 000000000..afe98f5eb --- /dev/null +++ b/dax/test/boltdb/helpers.go @@ -0,0 +1,50 @@ +package boltdb + +import ( + "os" + "testing" + + "github.com/molecula/featurebase/v3/dax/boltdb" + "github.com/stretchr/testify/assert" +) + +func MustGetDB(tb testing.TB) *boltdb.DB { + tb.Helper() + + f, err := os.CreateTemp("", "dax-boltdb") + assert.NoError(tb, err) + + dsn := "file:" + f.Name() + + db := boltdb.NewDB(dsn) + return db +} + +// MustOpenDB returns a new, open DB. Fatal on error. +func MustOpenDB(tb testing.TB) *boltdb.DB { + db := MustGetDB(tb) + + if err := db.Open(); err != nil { + tb.Fatal(err) + } + return db +} + +// MustCloseDB closes the DB. Fatal on error. +func MustCloseDB(tb testing.TB, db *boltdb.DB) { + tb.Helper() + if err := db.Close(); err != nil { + tb.Fatal(err) + } +} + +func CleanupDB(tb testing.TB, path string) { + tb.Helper() + + if path == "" { + return + } + if err := os.Remove(path); err != nil { + tb.Fatal(err) + } +} diff --git a/dax/test/datagen/container.go b/dax/test/datagen/container.go new file mode 100644 index 000000000..5ce45733e --- /dev/null +++ b/dax/test/datagen/container.go @@ -0,0 +1,80 @@ +package datagen + +import ( + "strconv" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/test/docker" +) + +var ImageName = docker.Getenv("DATAGEN_DOCKER_IMAGE", "dax/datagen") + +var ( + Seed = "123456" + Target = "mds" + Source = "custom" + StartFrom = 1 + EndAt = 1000 + BatchSize = 100 + PilosaIndex = "" + FeatureBaseOrganizationID = "" + FeatureBaseDatabaseID = "" + FeatureBaseTableName = "" + + MDSAddress = "mds:8080" +) + +const ( + NetworkName = "datagen-network" +) + +type Option func(*Container) + +func WithEndAt(endAt int) Option { + return func(c *Container) { + c.endAt = endAt + } +} + +type Container struct { + name string + mdsAddress dax.Address + endAt int +} + +func NewContainer(name string, mdsAddress dax.Address) *Container { + return &Container{ + name: name, + mdsAddress: mdsAddress, + } +} + +func (c *Container) Hostname() string { + return c.name +} + +func (c *Container) Env() map[string]string { + env := map[string]string{ + "GEN_MDS_ADDRESS": c.mdsAddress.String(), + "GEN_SEED": docker.Getenv("GEN_SEED", Seed), + "GEN_TARGET": Target, + "GEN_SOURCE": Source, + "GEN_START_FROM": strconv.Itoa(StartFrom), + "GEN_END_AT": strconv.Itoa(c.endAt), + "GEN_PILOSA_BATCH_SIZE": strconv.Itoa(BatchSize), + "GEN_FEATUREBASE_ORG_ID": docker.Getenv("GEN_FEATUREBASE_ORG_ID", FeatureBaseOrganizationID), + "GEN_FEATUREBASE_DB_ID": docker.Getenv("GEN_FEATUREBASE_DB_ID", FeatureBaseDatabaseID), + "GEN_FEATUREBASE_TABLE_NAME": docker.Getenv("GEN_FEATUREBASE_TABLE_NAME", FeatureBaseTableName), + "GEN_CUSTOM_CONFIG": docker.Getenv("GEN_CUSTOM_CONFIG", ""), + } + + return env +} + +func (c *Container) Cmd() []string { + return []string{"datagen"} +} + +func (c *Container) ExposedPorts() []string { + return []string{} +} diff --git a/dax/test/dax/dax_test.go b/dax/test/dax/dax_test.go new file mode 100644 index 000000000..33058b683 --- /dev/null +++ b/dax/test/dax/dax_test.go @@ -0,0 +1,1916 @@ +package dax + +import ( + "context" + "crypto/rand" + "encoding/json" + "errors" + "fmt" + "log" + "os" + "path/filepath" + "sort" + "strconv" + "strings" + "testing" + "time" + + fb "github.com/molecula/featurebase/v3" + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/mds/controller" + mdshttp "github.com/molecula/featurebase/v3/dax/mds/http" + queryerhttp "github.com/molecula/featurebase/v3/dax/queryer/http" + "github.com/molecula/featurebase/v3/dax/test" + "github.com/molecula/featurebase/v3/dax/test/datagen" + "github.com/molecula/featurebase/v3/dax/test/docker" + "github.com/molecula/featurebase/v3/dax/test/featurebase" + "github.com/molecula/featurebase/v3/dax/test/inspector" + "github.com/molecula/featurebase/v3/pql" + "github.com/molecula/featurebase/v3/sql3/parser" + planner_types "github.com/molecula/featurebase/v3/sql3/planner/types" + "github.com/molecula/featurebase/v3/sql3/test/defs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestDAXIntegration(t *testing.T) { + if testing.Short() { + t.Skip("skipping integration test") + } + + mdsStub := "mds" + computerStub := "computer" + queryerStub := "queryer" + writeloggerStub := "writelogger" + snapshotterStub := "snapshotter" + datagenStub := "datagen" + mydir, err := os.Getwd() + if err != nil { + t.Fatalf("getting CWD: %v", err) + } + coverVolume := docker.Volume{ + Type: "bind", + Source: filepath.Join(mydir, "../../../coverage-from-docker"), + Target: "/results", + } + + mdsNetworkName := "mds-network" + + addressFn := func(name string) dax.Address { + return dax.Address(name + ":8080") + } + + qual := dax.NewTableQualifier("acme", "db1") + + require := require.New(t) + + t.Run("SQL", func(t *testing.T) { + imagePull(t, featurebase.ImageName) + + fcName := uniqueContainerName(t, computerStub) + fc := featurebase.NewContainer(fcName, map[string]string{ + "FEATUREBASE_MDS_RUN": "true", + "FEATUREBASE_MDS_CONFIG_REGISTRATION_BATCH_TIMEOUT": "0s", + "FEATUREBASE_COMPUTER_RUN": "true", + "FEATUREBASE_QUERYER_RUN": "true", + }) + + dc := new(docker.Composer). + WithVolume(coverVolume, fc). + WithService(featurebase.ImageName, fc). + WithNetwork(uniqueNetworkName(t, featurebase.NetworkName)) + + defer dc.Down() + + // start featurebase.featurebase with no errors + require.NoError(dc.Up()) + waitForHealthy(t, fc, "", 10, time.Second) + waitForHealthy(t, fc, dax.ServicePrefixMDS, 10, time.Second) + waitForHealthy(t, fc, dax.ServicePrefixComputer, 10, time.Second) + waitForHealthy(t, fc, dax.ServicePrefixQueryer, 10, time.Second) + waitForStatus(t, fc, "NORMAL", 10, time.Second) + + tableTests := defs.TableTests + + // skips is a list of tests which are currently not passing in dax. We + // need to get these passing before alpha. + skips := []string{ + "testinsert/test-5", + "unoptestd/test-0", + "unoptestd/test-2", + "binoptesti_d/test-12", + "binoptestid_d/test-12", + "binoptestdec_i/test-12", + "binoptestdec_id/test-12", + "binoptestdec_d/test-12", + "table-82/test-3", + "table-82/test-8", + "cast_int/test-2", + "cast_int/test-7", + "cast_id/test-2", + "cast_string/test-12", + "cast_ts/test-7", + "sum_test/test-5", + "percentile_test/test-6", + "minmax_test/test-7", + "minmax_test/test-8", + "groupby_test/test-5", + "groupby_test/test-6", + "innerjointest/innerjoin-aggregate-groupby", + } + + doSkip := func(name string) bool { + for i := range skips { + if skips[i] == name { + return true + } + } + return false + } + + for i, test := range tableTests { + t.Run(test.Name(i), func(t *testing.T) { + + // Create a table with all field types. + if test.HasTable() { + sqlCheck(t, + fc, + addressFn(fcName), + qual, + test.CreateTable(), + `{"schema":{"fields":[]},"data":[],"error":"","warnings":null}`, + ) + } + + // Populate fields with data. + if test.HasTable() && test.HasData() { + sqlCheck(t, + fc, + addressFn(fcName), + qual, + test.InsertInto(t), + `{"schema":{"fields":[]},"data":[],"error":"","warnings":null}`, + ) + } + + for j, sqltest := range test.SQLTests { + t.Run(sqltest.Name(j), func(t *testing.T) { + if doSkip(test.Name(i) + "/" + sqltest.Name(j)) { + t.Skip("in dax skips list") + } + for _, sql := range sqltest.SQLs { + t.Run(fmt.Sprintf("sql-%s", sql), func(t *testing.T) { + log.Printf("SQL: %s", sql) + rows, headers, err := mustQueryRows(t, fc, addressFn(fcName), qual, sql, "") + + // Check expected error instead of results. + if sqltest.ExpErr != "" { + if assert.Error(t, err) { + assert.Contains(t, err.Error(), sqltest.ExpErr) + } + return + } + + require.NoError(err) + + // Check headers. + assert.ElementsMatch(t, sqltest.ExpHdrs, headers) + + // make a map of column name to header index + m := make(map[string]int) + for i := range headers { + m[headers[i].ColumnName] = i + } + + // Put the expRows in the same column order as the headers returned + // by the query. + exp := make([][]interface{}, len(sqltest.ExpRows)) + for i := range sqltest.ExpRows { + exp[i] = make([]interface{}, len(headers)) + for j := range sqltest.ExpHdrs { + targetIdx := m[sqltest.ExpHdrs[j].ColumnName] + assert.GreaterOrEqual(t, len(sqltest.ExpRows[i]), len(headers), + "expected row set has fewer columns than returned headers") + exp[i][targetIdx] = sqltest.ExpRows[i][j] + } + } + + if sqltest.SortStringKeys { + sortStringKeys(rows) + } + + switch sqltest.Compare { + case defs.CompareExactOrdered: + assert.Equal(t, len(sqltest.ExpRows), len(rows)) + assert.EqualValues(t, exp, rows) + case defs.CompareExactUnordered: + assert.Equal(t, len(sqltest.ExpRows), len(rows)) + assert.ElementsMatch(t, exp, rows) + case defs.CompareIncludedIn: + assert.Equal(t, sqltest.ExpRowCount, len(rows)) + for _, row := range rows { + assert.Contains(t, exp, row) + } + } + }) + } + }) + } + for j, pqltest := range test.PQLTests { + t.Run(pqltest.Name(j), func(t *testing.T) { + if doSkip(test.Name(i) + "/" + pqltest.Name(j)) { + t.Skip("in dax skips list") + } + for _, pql := range pqltest.PQLs { + t.Run(fmt.Sprintf("pql-%s", pql), func(t *testing.T) { + log.Printf("PQL: %s", pql) + rows, headers, err := mustQueryRows(t, fc, addressFn(fcName), qual, pql, pqltest.Table) + + // Check expected error instead of results. + if pqltest.ExpErr != "" { + if assert.Error(t, err) { + assert.Contains(t, err.Error(), pqltest.ExpErr) + } + return + } + + require.NoError(err) + + // Check headers. + assert.ElementsMatch(t, pqltest.ExpHdrs, headers) + + // make a map of column name to header index + m := make(map[string]int) + for i := range headers { + m[headers[i].ColumnName] = i + } + + // Put the expRows in the same column order as the headers returned + // by the query. + exp := make([][]interface{}, len(pqltest.ExpRows)) + for i := range pqltest.ExpRows { + exp[i] = make([]interface{}, len(headers)) + for j := range pqltest.ExpHdrs { + targetIdx := m[pqltest.ExpHdrs[j].ColumnName] + assert.GreaterOrEqual(t, len(pqltest.ExpRows[i]), len(headers), + "expected row set has fewer columns than returned headers") + exp[i][targetIdx] = pqltest.ExpRows[i][j] + } + } + + assert.Equal(t, len(pqltest.ExpRows), len(rows)) + assert.EqualValues(t, exp, rows) + + }) + } + }) + } + }) + } + + // time.Sleep(30000 * time.Second) + }) + + t.Run("Datagen", func(t *testing.T) { + imagePull(t, datagen.ImageName) + + // datagen + gc := datagen.NewContainer(uniqueContainerName(t, datagenStub), addressFn("")) + + dc := new(docker.Composer). + WithService(datagen.ImageName, gc). + WithNetwork(uniqueNetworkName(t, datagen.NetworkName)) + + defer dc.Down() + + t.Run("start datagen with no errors", func(t *testing.T) { + require.NoError(dc.Up()) + }) + }) + + t.Run("FeatureBase", func(t *testing.T) { + imagePull(t, featurebase.ImageName) + + fc := featurebase.NewContainer("base", nil) + + dc := new(docker.Composer). + WithVolume(coverVolume, fc). + WithService(featurebase.ImageName, fc). + WithNetwork(uniqueNetworkName(t, featurebase.NetworkName)) + + defer dc.Down() + + t.Run("start featurebase with no errors", func(t *testing.T) { + require.NoError(dc.Up()) + waitForHealthy(t, fc, "", 10, time.Second) + }) + }) + + t.Run("FeatureBase_MDS", func(t *testing.T) { + imagePull(t, featurebase.ImageName) + + fc := featurebase.NewContainer(uniqueContainerName(t, mdsStub), map[string]string{ + "FEATUREBASE_MDS_RUN": "true", + "FEATUREBASE_MDS_CONFIG_REGISTRATION_BATCH_TIMEOUT": "1ms", + }) + + dc := new(docker.Composer). + WithVolume(coverVolume, fc). + WithService(featurebase.ImageName, fc). + WithNetwork(uniqueNetworkName(t, featurebase.NetworkName)) + + defer dc.Down() + + t.Run("start featurebase.mds with no errors", func(t *testing.T) { + require.NoError(dc.Up()) + waitForHealthy(t, fc, "", 10, time.Second) + waitForHealthy(t, fc, dax.ServicePrefixMDS, 10, time.Second) + }) + }) + + // Queryer requires MDS to run. + t.Run("FeatureBase_Queryer_MDS", func(t *testing.T) { + imagePull(t, featurebase.ImageName) + + fc := featurebase.NewContainer(uniqueContainerName(t, queryerStub), map[string]string{ + "FEATUREBASE_MDS_RUN": "true", + "FEATUREBASE_MDS_CONFIG_REGISTRATION_BATCH_TIMEOUT": "1ms", + "FEATUREBASE_QUERYER_RUN": "true", + }) + + dc := new(docker.Composer). + WithVolume(coverVolume, fc). + WithService(featurebase.ImageName, fc). + WithNetwork(uniqueNetworkName(t, featurebase.NetworkName)) + + defer dc.Down() + + t.Run("start featurebase.queryer with no errors", func(t *testing.T) { + require.NoError(dc.Up()) + waitForHealthy(t, fc, "", 10, time.Second) + waitForHealthy(t, fc, dax.ServicePrefixQueryer, 10, time.Second) + }) + }) + + t.Run("FeatureBase_WriteLogger", func(t *testing.T) { + imagePull(t, featurebase.ImageName) + + fc := featurebase.NewContainer(uniqueContainerName(t, writeloggerStub), map[string]string{ + "FEATUREBASE_WRITELOGGER_RUN": "true", + }) + + dc := new(docker.Composer). + WithVolume(coverVolume, fc). + WithService(featurebase.ImageName, fc). + WithNetwork(uniqueNetworkName(t, featurebase.NetworkName)) + + defer dc.Down() + + t.Run("start featurebase.writelogger with no errors", func(t *testing.T) { + require.NoError(dc.Up()) + waitForHealthy(t, fc, "", 10, time.Second) + waitForHealthy(t, fc, dax.ServicePrefixWriteLogger, 10, time.Second) + }) + }) + + t.Run("FeatureBase_Snapshotter", func(t *testing.T) { + imagePull(t, featurebase.ImageName) + + fc := featurebase.NewContainer(uniqueContainerName(t, snapshotterStub), map[string]string{ + "FEATUREBASE_SNAPSHOTTER_RUN": "true", + }) + + dc := new(docker.Composer). + WithVolume(coverVolume, fc). + WithService(featurebase.ImageName, fc). + WithNetwork(uniqueNetworkName(t, featurebase.NetworkName)) + + defer dc.Down() + + t.Run("start featurebase.snapshotter with no errors", func(t *testing.T) { + require.NoError(dc.Up()) + waitForHealthy(t, fc, "", 10, time.Second) + waitForHealthy(t, fc, dax.ServicePrefixSnapshotter, 10, time.Second) + }) + }) + + // Computer requires MDS to run. + t.Run("FeatureBase_Computer", func(t *testing.T) { + imagePull(t, featurebase.ImageName) + + fcName := uniqueContainerName(t, computerStub) + fc := featurebase.NewContainer(fcName, map[string]string{ + "FEATUREBASE_MDS_RUN": "true", + "FEATUREBASE_MDS_CONFIG_REGISTRATION_BATCH_TIMEOUT": "1ms", + "FEATUREBASE_COMPUTER_RUN": "true", + }) + + dc := new(docker.Composer). + WithVolume(coverVolume, fc). + WithService(featurebase.ImageName, fc). + WithNetwork(uniqueNetworkName(t, featurebase.NetworkName)) + + defer dc.Down() + + t.Run("start featurebase.featurebase with no errors", func(t *testing.T) { + require.NoError(dc.Up()) + waitForHealthy(t, fc, "", 10, time.Second) + waitForHealthy(t, fc, dax.ServicePrefixComputer, 10, time.Second) + waitForStatus(t, fc, "NORMAL", 10, time.Second) + }) + + t.Run("check the schema endpoint", func(t *testing.T) { + inspect(t, + fc, + addressFn(fcName), + "computer/schema", + nil, + getExpFEqual(`{"indexes":[]}`), + ) + }) + }) + + t.Run("MDS_FeatureBase", func(t *testing.T) { + imagePull(t, featurebase.ImageName) + + mcName := uniqueContainerName(t, mdsStub) + mc := featurebase.NewContainer(mcName, map[string]string{ + "FEATUREBASE_MDS_RUN": "true", + "FEATUREBASE_MDS_CONFIG_REGISTRATION_BATCH_TIMEOUT": "1ms", + }) + + fcName := uniqueContainerName(t, computerStub) + fc := featurebase.NewContainer(fcName, map[string]string{ + "FEATUREBASE_COMPUTER_RUN": "true", + "FEATUREBASE_COMPUTER_CONFIG_MDS_ADDRESS": string(addressFn(mcName)), + }) + + tableName := dax.TableName("tbl") + keyedTbl := test.TestQualifiedTableWithID(t, qual, "", tableName, 12, true) + expTbl := test.TestQualifiedTableWithID(t, qual, "someid", tableName, 12, true) + + dc := new(docker.Composer). + WithVolume(coverVolume, fc). + WithService(featurebase.ImageName, mc). + WithService(featurebase.ImageName, fc). + WithNetwork(uniqueNetworkName(t, mdsNetworkName)) + + defer dc.Down() + + t.Run("start the containers with no errors", func(t *testing.T) { + require.NoError(dc.Up()) + waitForStatus(t, fc, "NORMAL", 10, time.Second) + }) + + t.Run("add a keyed table", func(t *testing.T) { + inspect(t, + mc, + addressFn(mcName), + "mds/create-table", + keyedTbl, + getExpFTable(expTbl), + ) + }) + + t.Run("check the translate-nodes endpoint", func(t *testing.T) { + data := mdshttp.TranslateNodesRequest{ + Table: keyedTbl.QualifiedID(), + Partitions: dax.PartitionNums{3, 5, 8}, + } + + inspect(t, + mc, + addressFn(mcName), + "mds/translate-nodes", + data, + getExpFTranslateResponse(mdshttp.TranslateNodesResponse{ + TranslateNodes: []controller.TranslateNode{ + { + Address: addressFn(fcName), + Partitions: dax.PartitionNums{3, 5, 8}, + }, + }}), + ) + }) + + // The steps above result in featurebase receiving a directive which + // creates the table. + t.Run("check the schema endpoint", func(t *testing.T) { + inspect(t, + fc, + addressFn(fcName), + "computer/schema", + nil, + getExpFContains(`{"indexes":[{"name":"`, `"options":{"keys":true,"trackExistence":true,"partitionN":0,"description":""},"fields":[],"shardWidth":1048576}`), + ) + }) + }) + + t.Run("MDS_FeatureBase_Datagen_Queryer", func(t *testing.T) { + imagePull(t, featurebase.ImageName) + imagePull(t, datagen.ImageName) + + tableName := dax.TableName("tbl") + partitionN := 256 + tbl := dax.NewTable(tableName) + tbl.PartitionN = partitionN + tbl.Fields = []*dax.Field{ + {Name: "_id", Type: "id"}, + {Name: "an_int", Type: "int", + Options: dax.FieldOptions{ + Min: pql.NewDecimal(0, 0), + Max: pql.NewDecimal(500, 0), + }, + }, + {Name: "a_random_string", Type: "string"}, + {Name: "an_id_set", Type: "idset"}, + {Name: "a_string_set", Type: "stringset"}, + } + qtbl := dax.NewQualifiedTable(qual, tbl) + + mcName := uniqueContainerName(t, mdsStub) + mc := featurebase.NewContainer(mcName, map[string]string{ + "FEATUREBASE_MDS_RUN": "true", + "FEATUREBASE_MDS_CONFIG_REGISTRATION_BATCH_TIMEOUT": "1ms", + }) + + fc := featurebase.NewContainer(uniqueContainerName(t, computerStub), map[string]string{ + "FEATUREBASE_COMPUTER_RUN": "true", + "FEATUREBASE_COMPUTER_CONFIG_MDS_ADDRESS": string(addressFn(mcName)), + }) + + qcName := uniqueContainerName(t, queryerStub) + qc := featurebase.NewContainer(qcName, map[string]string{ + "FEATUREBASE_QUERYER_RUN": "true", + "FEATUREBASE_QUERYER_CONFIG_MDS_ADDRESS": string(addressFn(mcName)), + }) + + networkName := uniqueNetworkName(t, mdsNetworkName) + + dc := new(docker.Composer). + WithVolume(coverVolume, fc, mc, qc). + WithService(featurebase.ImageName, mc). + WithService(featurebase.ImageName, fc). + WithService(featurebase.ImageName, qc). + WithNetwork(networkName) + + t.Run("start the containers with no errors", func(t *testing.T) { + require.NoError(dc.Up()) + waitForHealthy(t, mc, dax.ServicePrefixMDS, 10, time.Second) + waitForHealthy(t, qc, dax.ServicePrefixQueryer, 10, time.Second) + waitForHealthy(t, fc, dax.ServicePrefixComputer, 10, time.Second) + waitForStatus(t, fc, "NORMAL", 10, time.Second) + }) + defer dc.Down() + + t.Run("add a table", func(t *testing.T) { + sqlCheck(t, + mc, + addressFn(qcName), + qtbl.Qualifier(), + qtbl.CreateSQL(), + `{"schema":{"fields":[]},"data":[],"error":"","warnings":null}`, + ) + }) + + t.Run("ingest some data", func(t *testing.T) { + // datagen + assert.NoError(t, docker.Setenv("GEN_FEATUREBASE_ORG_ID", string(qual.OrganizationID))) + assert.NoError(t, docker.Setenv("GEN_FEATUREBASE_DB_ID", string(qual.DatabaseID))) + assert.NoError(t, docker.Setenv("GEN_FEATUREBASE_TABLE_NAME", string(qtbl.Name))) + assert.NoError(t, docker.Setenv("GEN_CUSTOM_CONFIG", "/testdata/basic.yaml")) + assert.NoError(t, docker.Setenv("GEN_USE_SHARD_TRANSACTIONAL_ENDPOINT", "true")) + gcName := uniqueContainerName(t, datagenStub) + gc := datagen.NewContainer(gcName, addressFn(mcName)) + _, err := docker.ContainerCreate(datagen.ImageName, gc, []docker.Volume{}) + assert.NoError(t, err) + assert.NoError(t, docker.NetworkConnect(networkName, gcName)) + assert.NoError(t, docker.ContainerStartWithLogging(gcName)) + defer func() { + docker.ContainerStop(gcName) + docker.ContainerRemove(gcName) + }() + + // Wait for datagen to finish ingest. + require.NoError(docker.ContainerWait(gc.Hostname())) + }) + + t.Run("query some data", func(t *testing.T) { + // PQL: `Count(Row(a_string_set="A90B"))`, + sqlCheck(t, + mc, + addressFn(qcName), + qual, + `select count(*) as cnt from tbl where setcontains(a_string_set, 'A90B')`, + `{"schema":{"fields":[{"name":"cnt","type":"INT"}]},"data":[[1]],"error":"","warnings":null}`, + ) + }) + + // time.Sleep(30000 * time.Second) + }) + + t.Run("MDS_Poller_FeatureBase", func(t *testing.T) { + imagePull(t, featurebase.ImageName) + + computerName1 := uniqueContainerName(t, computerStub+"1") + computerName2 := uniqueContainerName(t, computerStub+"2") + + tableName := dax.TableName("tbl") + partitionN := 12 + qtbl := test.TestQualifiedTableWithID(t, qual, "", tableName, partitionN, true) + expTbl := test.TestQualifiedTableWithID(t, qual, "someID", tableName, partitionN, true) + + mcName := uniqueContainerName(t, mdsStub) + mc := featurebase.NewContainer(mcName, map[string]string{ + "FEATUREBASE_MDS_RUN": "true", + "FEATUREBASE_MDS_CONFIG_REGISTRATION_BATCH_TIMEOUT": "1ms", + }) + + env := map[string]string{ + "FEATUREBASE_COMPUTER_RUN": "true", + "FEATUREBASE_COMPUTER_CONFIG_MDS_ADDRESS": string(addressFn(mcName)), + } + fc1 := featurebase.NewContainer(computerName1, env) + fc2 := featurebase.NewContainer(computerName2, env) + + dc := new(docker.Composer). + WithVolume(coverVolume, fc1, fc2, mc). + WithService(featurebase.ImageName, mc). + WithService(featurebase.ImageName, fc1). + WithService(featurebase.ImageName, fc2). + WithNetwork(uniqueNetworkName(t, mdsNetworkName)) + + t.Run("start the containers with no errors", func(t *testing.T) { + require.NoError(dc.Up()) + waitForHealthy(t, mc, dax.ServicePrefixMDS, 10, time.Second) + waitForHealthy(t, fc1, dax.ServicePrefixComputer, 10, time.Second) + waitForStatus(t, fc1, "NORMAL", 10, time.Second) + waitForHealthy(t, fc2, dax.ServicePrefixComputer, 10, time.Second) + waitForStatus(t, fc2, "NORMAL", 10, time.Second) + }) + defer dc.Down() + + t.Run("add a keyed table", func(t *testing.T) { + inspect(t, + mc, + addressFn(mcName), + "mds/create-table", + qtbl, + getExpFTable(expTbl), + ) + }) + + partitions := dax.PartitionNums{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11} + + // ensure partitions are covered + t.Run("check the translate-nodes endpoint", func(t *testing.T) { + data := mdshttp.TranslateNodesRequest{ + Table: qtbl.QualifiedID(), + Partitions: partitions, + IsWrite: true, + } + + inspect(t, + mc, + addressFn(mcName), + "mds/translate-nodes", + data, + getExpFTranslateResponse(mdshttp.TranslateNodesResponse{ + TranslateNodes: []controller.TranslateNode{ + { + Address: addressFn(computerName1), + Partitions: dax.PartitionNums{0, 2, 4, 6, 8, 10}, + }, + { + Address: addressFn(computerName2), + Partitions: dax.PartitionNums{1, 3, 5, 7, 9, 11}, + }, + }}), + ) + }) + + // stop featurebase 1 (may need to sleep) + t.Run("stop a container", func(t *testing.T) { + docker.ContainerStop(computerName1) + }) + + // Give the poller time to recognize the node is gone. + // TODO: implement this without a sleep. + time.Sleep(20 * time.Second) + + // ensure paritions are still covered + t.Run("check the translate-nodes endpoint again", func(t *testing.T) { + data := mdshttp.TranslateNodesRequest{ + Table: qtbl.QualifiedID(), + Partitions: partitions, + } + + inspect(t, + mc, + addressFn(mcName), + "mds/translate-nodes", + data, + getExpFTranslateResponse(mdshttp.TranslateNodesResponse{ + TranslateNodes: []controller.TranslateNode{ + { + Address: addressFn(computerName2), + Partitions: dax.PartitionNums{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}, + }, + }}), + ) + }) + + //time.Sleep(3000 * time.Second) + }) + + t.Run("Node_Recovery", func(t *testing.T) { + imagePull(t, featurebase.ImageName) + + computerName1 := uniqueContainerName(t, computerStub+"1") + computerName2 := uniqueContainerName(t, computerStub+"2") + + tableName := dax.TableName("tbl") + partitionN := 12 + + tbl := dax.NewTable(tableName) + tbl.PartitionN = partitionN + tbl.Fields = []*dax.Field{ + {Name: "_id", Type: "id"}, + {Name: "an_int", Type: "int", + Options: dax.FieldOptions{ + Min: pql.NewDecimal(0, 0), + Max: pql.NewDecimal(500, 0), + }, + }, + {Name: "a_random_string", Type: "string"}, + {Name: "an_id_set", Type: "idset"}, + {Name: "a_string_set", Type: "stringset"}, + } + qtbl := dax.NewQualifiedTable(qual, tbl) + + mcName := uniqueContainerName(t, mdsStub) + mc := featurebase.NewContainer(mcName, map[string]string{ + "FEATUREBASE_MDS_RUN": "true", + "FEATUREBASE_MDS_CONFIG_REGISTRATION_BATCH_TIMEOUT": "1ms", + }) + + wcName := uniqueContainerName(t, writeloggerStub) + wc := featurebase.NewContainer(wcName, map[string]string{ + "FEATUREBASE_WRITELOGGER_RUN": "true", + "FEATUREBASE_WRITELOGGER_CONFIG_DATA_DIR": "/tmp", + }) + + fc1 := featurebase.NewContainer(computerName1, map[string]string{ + "FEATUREBASE_COMPUTER_RUN": "true", + "FEATUREBASE_COMPUTER_CONFIG_MDS_ADDRESS": string(addressFn(mcName)), + "FEATUREBASE_COMPUTER_CONFIG_WRITE_LOGGER": string(addressFn(wcName)), + }) + + qcName := uniqueContainerName(t, queryerStub) + qc := featurebase.NewContainer(qcName, map[string]string{ + "FEATUREBASE_QUERYER_RUN": "true", + "FEATUREBASE_QUERYER_CONFIG_MDS_ADDRESS": string(addressFn(mcName)), + }) + + networkName := uniqueNetworkName(t, mdsNetworkName) + + dc := new(docker.Composer). + WithVolume(coverVolume, mc, wc, qc, fc1). + WithService(featurebase.ImageName, mc). + WithService(featurebase.ImageName, wc). + WithService(featurebase.ImageName, qc). + WithService(featurebase.ImageName, fc1). + WithNetwork(networkName) + + t.Run("start the containers with no errors", func(t *testing.T) { + require.NoError(dc.Up()) + waitForHealthy(t, mc, dax.ServicePrefixMDS, 10, time.Second) + waitForHealthy(t, wc, dax.ServicePrefixWriteLogger, 10, time.Second) + waitForHealthy(t, qc, dax.ServicePrefixQueryer, 10, time.Second) + waitForHealthy(t, fc1, dax.ServicePrefixComputer, 10, time.Second) + waitForStatus(t, fc1, "NORMAL", 10, time.Second) + }) + defer dc.Down() + + t.Run("add a table", func(t *testing.T) { + sqlCheck(t, + mc, + addressFn(qcName), + qtbl.Qualifier(), + qtbl.CreateSQL(), + `{"schema":{"fields":[]},"data":[],"error":"","warnings":null}`, + ) + }) + + t.Run("ingest some data", func(t *testing.T) { + // datagen + assert.NoError(t, docker.Setenv("GEN_FEATUREBASE_ORG_ID", string(qual.OrganizationID))) + assert.NoError(t, docker.Setenv("GEN_FEATUREBASE_DB_ID", string(qual.DatabaseID))) + assert.NoError(t, docker.Setenv("GEN_FEATUREBASE_TABLE_NAME", string(qtbl.Name))) + assert.NoError(t, docker.Setenv("GEN_CUSTOM_CONFIG", "/testdata/basic.yaml")) + assert.NoError(t, docker.Setenv("GEN_USE_SHARD_TRANSACTIONAL_ENDPOINT", "true")) + gcName := uniqueContainerName(t, datagenStub) + gc := datagen.NewContainer(gcName, addressFn(mcName)) + _, err := docker.ContainerCreate(datagen.ImageName, gc, []docker.Volume{}) + assert.NoError(t, err) + assert.NoError(t, docker.NetworkConnect(networkName, gcName)) + assert.NoError(t, docker.ContainerStartWithLogging(gcName)) + defer func() { + docker.ContainerStop(gcName) + docker.ContainerRemove(gcName) + }() + + // Wait for datagen to finish ingest. + require.NoError(docker.ContainerWait(gc.Hostname())) + }) + + t.Run("query some data", func(t *testing.T) { + // PQL: `Count(Row(a_string_set="A90B"))`, + sqlCheck(t, + mc, + addressFn(qcName), + qual, + `select count(*) as cnt from tbl where setcontains(a_string_set, 'A90B')`, + `{"schema":{"fields":[{"name":"cnt","type":"INT"}]},"data":[[1]],"error":"","warnings":null}`, + ) + }) + + // stop featurebase 1 (may need to sleep) + t.Run("stop a container", func(t *testing.T) { + docker.ContainerStop(computerName1) + }) + + // Give the poller time to recognize the node is gone. + // TODO: implement this without a sleep. + time.Sleep(25 * time.Second) + + t.Run("start a new compute node", func(t *testing.T) { + fc2 := featurebase.NewContainer(computerName2, map[string]string{ + "FEATUREBASE_COMPUTER_RUN": "true", + "FEATUREBASE_COMPUTER_CONFIG_MDS_ADDRESS": string(addressFn(mcName)), + "FEATUREBASE_COMPUTER_CONFIG_WRITE_LOGGER": string(addressFn(wcName)), + }) + _, err := docker.ContainerCreate(featurebase.ImageName, fc2, []docker.Volume{coverVolume}) + assert.NoError(t, err) + assert.NoError(t, docker.NetworkConnect(networkName, computerName2)) + assert.NoError(t, docker.ContainerStart(computerName2)) + + waitForHealthy(t, fc2, dax.ServicePrefixComputer, 10, time.Second) + waitForStatus(t, fc2, "NORMAL", 10, time.Second) + }) + // This closes the container created in "start a new compute node", but + // this should really only be called if that sub-test runs successfully. + defer func() { + docker.ContainerStop(computerName2) + docker.ContainerRemove(computerName2) + }() + + t.Run("query the same data", func(t *testing.T) { + // PQL: `Count(Row(a_string_set="A90B"))`, + sqlCheck(t, + mc, + addressFn(qcName), + qual, + `select count(*) as cnt from tbl where setcontains(a_string_set, 'A90B')`, + `{"schema":{"fields":[{"name":"cnt","type":"INT"}]},"data":[[1]],"error":"","warnings":null}`, + ) + }) + + //time.Sleep(3000 * time.Second) + }) + + t.Run("Node_Recovery_Snapshot", func(t *testing.T) { + imagePull(t, featurebase.ImageName) + + computerName1 := uniqueContainerName(t, computerStub+"1") + computerName2 := uniqueContainerName(t, computerStub+"2") + + tableName := dax.TableName("tbl") + partitionN := 12 + + tbl := dax.NewTable(tableName) + tbl.PartitionN = partitionN + tbl.Fields = []*dax.Field{ + {Name: "_id", Type: "id"}, + {Name: "an_int", Type: "int", + Options: dax.FieldOptions{ + Min: pql.NewDecimal(0, 0), + Max: pql.NewDecimal(500, 0), + }, + }, + {Name: "a_random_string", Type: "string"}, + {Name: "an_id_set", Type: "idset"}, + {Name: "a_string_set", Type: "stringset"}, + } + qtbl := dax.NewQualifiedTable(qual, tbl) + + mcName := uniqueContainerName(t, mdsStub) + mc := featurebase.NewContainer(mcName, map[string]string{ + "FEATUREBASE_MDS_RUN": "true", + "FEATUREBASE_MDS_CONFIG_REGISTRATION_BATCH_TIMEOUT": "1ms", + }) + + wcName := uniqueContainerName(t, writeloggerStub) + wc := featurebase.NewContainer(wcName, map[string]string{ + "FEATUREBASE_WRITELOGGER_RUN": "true", + "FEATUREBASE_WRITELOGGER_CONFIG_DATA_DIR": "/tmp", + }) + + scName := uniqueContainerName(t, snapshotterStub) + sc := featurebase.NewContainer(scName, map[string]string{ + "FEATUREBASE_SNAPSHOTTER_RUN": "true", + "FEATUREBASE_SNAPSHOTTER_CONFIG_DATA_DIR": "/tmp", + }) + + fc1 := featurebase.NewContainer(computerName1, map[string]string{ + "FEATUREBASE_COMPUTER_RUN": "true", + "FEATUREBASE_COMPUTER_CONFIG_MDS_ADDRESS": string(addressFn(mcName)), + "FEATUREBASE_COMPUTER_CONFIG_WRITE_LOGGER": string(addressFn(wcName)), + "FEATUREBASE_COMPUTER_CONFIG_SNAPSHOTTER": string(addressFn(scName)), + }) + + qcName := uniqueContainerName(t, queryerStub) + qc := featurebase.NewContainer(qcName, map[string]string{ + "FEATUREBASE_QUERYER_RUN": "true", + "FEATUREBASE_QUERYER_CONFIG_MDS_ADDRESS": string(addressFn(mcName)), + }) + + networkName := uniqueNetworkName(t, mdsNetworkName) + + dc := new(docker.Composer). + WithVolume(coverVolume, mc, wc, sc, qc, fc1). + WithService(featurebase.ImageName, mc). + WithService(featurebase.ImageName, wc). + WithService(featurebase.ImageName, sc). + WithService(featurebase.ImageName, qc). + WithService(featurebase.ImageName, fc1). + WithNetwork(networkName) + + t.Run("start the containers with no errors", func(t *testing.T) { + require.NoError(dc.Up()) + waitForHealthy(t, mc, dax.ServicePrefixMDS, 10, time.Second) + waitForHealthy(t, wc, dax.ServicePrefixWriteLogger, 10, time.Second) + waitForHealthy(t, sc, dax.ServicePrefixSnapshotter, 10, time.Second) + waitForHealthy(t, qc, dax.ServicePrefixQueryer, 10, time.Second) + waitForHealthy(t, fc1, dax.ServicePrefixComputer, 10, time.Second) + waitForStatus(t, fc1, "NORMAL", 10, time.Second) + }) + defer dc.Down() + + // TODO(tlt): using SQL to create the table doesn't work in tests + // because we don't know the TableID that gets auto-generated. And + // therefore we can't use that TableID later to, for example, trigger a + // snapshot. For now we'll go directly to MDS to create the table (where + // we can specify the TableID here in the test), but we probably need a + // way to lookup a TableID given qual/table-name. + // + // t.Run("add a non-keyed table", func(t *testing.T) { + // sqlCheck(t, + // mc, + // queryerAddress, + // qtbl.Qualifier(), + // qtbl.CreateSQL(), + // `{"schema":{"fields":[]},"data":[],"error":"","warnings":null}`, + // ) + // }) + t.Run("add a non-keyed table", func(t *testing.T) { + inspect(t, + mc, + addressFn(mcName), + "mds/create-table", + qtbl, + getExpFContains( + `"name":"tbl","fields":[`, + `{"name":"_id","type":"id","options":{"min":0,"max":0,"epoch":"0001-01-01T00:00:00Z"}}`, + `{"name":"an_int","type":"int","options":{"min":0,"max":500,"epoch":"0001-01-01T00:00:00Z"}}`, + `{"name":"a_random_string","type":"string","options":{"min":0,"max":0,"epoch":"0001-01-01T00:00:00Z"}}`, + `{"name":"an_id_set","type":"idset","options":{"min":0,"max":0,"epoch":"0001-01-01T00:00:00Z"}}`, + `{"name":"a_string_set","type":"stringset","options":{"min":0,"max":0,"epoch":"0001-01-01T00:00:00Z"}}`, + `"partitionN":12,"org-id":"acme","db-id":"db1"}`, + ), + ) + }) + + t.Run("ingest some data", func(t *testing.T) { + // datagen + assert.NoError(t, docker.Setenv("GEN_FEATUREBASE_ORG_ID", string(qual.OrganizationID))) + assert.NoError(t, docker.Setenv("GEN_FEATUREBASE_DB_ID", string(qual.DatabaseID))) + assert.NoError(t, docker.Setenv("GEN_FEATUREBASE_TABLE_NAME", string(qtbl.Name))) + //assert.NoError(t, docker.Setenv("GEN_SEED", "1")) + assert.NoError(t, docker.Setenv("GEN_CUSTOM_CONFIG", "/testdata/basic.yaml")) + assert.NoError(t, docker.Setenv("GEN_USE_SHARD_TRANSACTIONAL_ENDPOINT", "true")) + gcName := uniqueContainerName(t, datagenStub) + gc := datagen.NewContainer(gcName, addressFn(mcName)) + _, err := docker.ContainerCreate(datagen.ImageName, gc, []docker.Volume{}) + assert.NoError(t, err) + assert.NoError(t, docker.NetworkConnect(networkName, gcName)) + assert.NoError(t, docker.ContainerStartWithLogging(gcName)) + defer func() { + docker.ContainerStop(gcName) + docker.ContainerRemove(gcName) + }() + + // Wait for datagen to finish ingest. + require.NoError(docker.ContainerWait(gc.Hostname())) + }) + + t.Run("query some data", func(t *testing.T) { + // PQL: `Count(Row(a_string_set="B25A"))`, + sqlCheck(t, + mc, + addressFn(qcName), + qual, + `select count(*) as cnt from tbl where setcontains(a_string_set, 'B25A')`, + `{"schema":{"fields":[{"name":"cnt","type":"INT"}]},"data":[[1]],"error":"","warnings":null}`, + ) + }) + + // snapshot shards and partitions + t.Run("snapshot shards", func(t *testing.T) { + data := mdshttp.SnapshotShardRequest{ + Table: qtbl.QualifiedID(), + Shard: 0, + } + inspect(t, + mc, + addressFn(mcName), + "mds/snapshot/shard-data", + data, + getExpFEqual(``), + ) + }) + t.Run("snapshot partitions", func(t *testing.T) { + data := mdshttp.SnapshotFieldKeysRequest{ + Table: qtbl.QualifiedID(), + Field: "a_string_set", + } + inspect(t, + mc, + addressFn(mcName), + "mds/snapshot/field-keys", + data, + getExpFEqual(``), + ) + }) + + // ingest some more data (seed 1) + t.Run("ingest some more data", func(t *testing.T) { + // datagen + assert.NoError(t, docker.Setenv("GEN_FEATUREBASE_ORG_ID", string(qual.OrganizationID))) + assert.NoError(t, docker.Setenv("GEN_FEATUREBASE_DB_ID", string(qual.DatabaseID))) + assert.NoError(t, docker.Setenv("GEN_FEATUREBASE_TABLE_NAME", string(qtbl.Name))) + assert.NoError(t, docker.Setenv("GEN_SEED", "1")) + assert.NoError(t, docker.Setenv("GEN_CUSTOM_CONFIG", "/testdata/basic.yaml")) + assert.NoError(t, docker.Setenv("GEN_USE_SHARD_TRANSACTIONAL_ENDPOINT", "true")) + gcName := uniqueContainerName(t, datagenStub) + gc := datagen.NewContainer(gcName, addressFn(mcName)) + _, err := docker.ContainerCreate(datagen.ImageName, gc, []docker.Volume{}) + assert.NoError(t, err) + assert.NoError(t, docker.NetworkConnect(networkName, gcName)) + assert.NoError(t, docker.ContainerStartWithLogging(gcName)) + defer func() { + docker.ContainerStop(gcName) + docker.ContainerRemove(gcName) + }() + + // Wait for datagen to finish ingest. + require.NoError(docker.ContainerWait(gc.Hostname())) + }) + + t.Run("query the new data", func(t *testing.T) { + // PQL: `Count(Row(a_string_set="B25A"))`, + sqlCheck(t, + mc, + addressFn(qcName), + qual, + `select count(*) as cnt from tbl where setcontains(a_string_set, 'B25A')`, + `{"schema":{"fields":[{"name":"cnt","type":"INT"}]},"data":[[2]],"error":"","warnings":null}`, + ) + }) + + // stop featurebase 1 (may need to sleep) + t.Run("stop a container", func(t *testing.T) { + docker.ContainerStop(computerName1) + }) + + // Give the poller time to recognize the node is gone. + // TODO: implement this without a sleep. + time.Sleep(25 * time.Second) + + t.Run("start a new compute node", func(t *testing.T) { + fc2 := featurebase.NewContainer(computerName2, map[string]string{ + "FEATUREBASE_COMPUTER_RUN": "true", + "FEATUREBASE_COMPUTER_CONFIG_MDS_ADDRESS": string(addressFn(mcName)), + "FEATUREBASE_COMPUTER_CONFIG_WRITE_LOGGER": string(addressFn(wcName)), + "FEATUREBASE_COMPUTER_CONFIG_SNAPSHOTTER": string(addressFn(scName)), + }) + _, err := docker.ContainerCreate(featurebase.ImageName, fc2, []docker.Volume{coverVolume}) + assert.NoError(t, err) + assert.NoError(t, docker.NetworkConnect(networkName, computerName2)) + assert.NoError(t, docker.ContainerStart(computerName2)) + + waitForHealthy(t, fc2, dax.ServicePrefixComputer, 10, time.Second) + waitForStatus(t, fc2, "NORMAL", 10, time.Second) + }) + // This closes the container created in "start a new compute node", but + // this should really only be called if that sub-test runs successfully. + defer func() { + docker.ContainerStop(computerName2) + docker.ContainerRemove(computerName2) + }() + + t.Run("query the same data", func(t *testing.T) { + // PQL: `Count(Row(a_string_set="B25A"))`, + sqlCheck(t, + mc, + addressFn(qcName), + qual, + `select count(*) as cnt from tbl where setcontains(a_string_set, 'B25A')`, + `{"schema":{"fields":[{"name":"cnt","type":"INT"}]},"data":[[2]],"error":"","warnings":null}`, + ) + }) + + // TODO(tlt): it doesn't really make sense to have this tacked on to the + // end of this test. There should be a separate test run which calls + // this instead of the individual snapshot calls (in this test, above). + // What we really need to do is create a test which can + // create/snapshot/restore different table types (keyed/non-keyed), and + // test different snapshot methods (individual shards, whole table, + // etc.). + t.Run("snapshot table", func(t *testing.T) { + data := qtbl.QualifiedID() + inspect(t, + mc, + addressFn(mcName), + "mds/snapshot", + data, + getExpFEqual(``), + ) + }) + + }) + + t.Run("MDS_Persistence", func(t *testing.T) { + imagePull(t, featurebase.ImageName) + + computerName1 := uniqueContainerName(t, computerStub+"1") + computerName2 := uniqueContainerName(t, computerStub+"2") + + // Note: this test is able to bring up a new mds container with a + // different name because it doesn't rely on any other container + // maintaining a link to a reliable, constant "mds" address (i.e. this + // test doesn't use queryer, or rely on a compute node registering with + // mdsAddress2). We could probably share the same name for both + // containers, and therefore not prevent other nodes from communicating + // with the new mds container, but the purpose of this test is to + // explicity show that the data is persistent across two, distinct mds + // containers. + mdsName1 := uniqueContainerName(t, mdsStub+"1") + mdsName2 := uniqueContainerName(t, mdsStub+"2") + + mdsAddress1 := addressFn(mdsName1) + mdsAddress2 := addressFn(mdsName2) + + tableName := dax.TableName("tbl") + partitionN := 12 + qtbl := test.TestQualifiedTableWithID(t, qual, "", tableName, partitionN, true) + + mdsVolumeSource := "mds-storage" + mdsVolumeTarget := "/storage" + dsn := "file:" + mdsVolumeTarget + "/mds.boltdb" + + mc1 := featurebase.NewContainer(mdsName1, map[string]string{ + "FEATUREBASE_MDS_RUN": "true", + "FEATUREBASE_MDS_CONFIG_REGISTRATION_BATCH_TIMEOUT": "1ms", + "FEATUREBASE_STORAGE_METHOD": "boltdb", + "FEATUREBASE_STORAGE_DSN": dsn, + }) + var mc2 *featurebase.Container + + env := map[string]string{ + "FEATUREBASE_COMPUTER_RUN": "true", + "FEATUREBASE_COMPUTER_CONFIG_MDS_ADDRESS": string(mdsAddress1), + } + fc1 := featurebase.NewContainer(computerName1, env) + + mdsVolume := docker.Volume{ + Source: mdsVolumeSource, + Target: mdsVolumeTarget, + } + + networkName := uniqueNetworkName(t, mdsNetworkName) + + dc := new(docker.Composer). + WithVolume(mdsVolume, mc1). + WithVolume(coverVolume, mc1, fc1). + WithService(featurebase.ImageName, mc1). + WithService(featurebase.ImageName, fc1). + WithNetwork(networkName) + + t.Run("start the containers with no errors", func(t *testing.T) { + require.NoError(dc.Up()) + waitForHealthy(t, mc1, dax.ServicePrefixMDS, 10, time.Second) + waitForHealthy(t, fc1, dax.ServicePrefixComputer, 10, time.Second) + waitForStatus(t, fc1, "NORMAL", 10, time.Second) + }) + defer dc.Down() + + t.Run("add a keyed table", func(t *testing.T) { + inspect(t, + mc1, + mdsAddress1, + "mds/create-table", + qtbl, + getExpFContains(`"name":"tbl","fields":[{"name":"_id","type":"string","options":{"min":0,"max":0,"epoch":"0001-01-01T00:00:00Z"}}],"partitionN":12,"org-id":"acme","db-id":"db1"}`), + ) + }) + + partitions := dax.PartitionNums{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11} + + // ensure partitions are covered + t.Run("check the translate-nodes endpoint", func(t *testing.T) { + data := mdshttp.TranslateNodesRequest{ + Table: qtbl.QualifiedID(), + Partitions: partitions, + IsWrite: true, + } + + inspect(t, + mc1, + mdsAddress1, + "mds/translate-nodes", + data, + getExpFTranslateResponse(mdshttp.TranslateNodesResponse{ + TranslateNodes: []controller.TranslateNode{ + { + Address: addressFn(computerName1), + Partitions: dax.PartitionNums{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}, + }, + }}), + ) + }) + + // stop mds (may need to sleep) + t.Run("stop mds container", func(t *testing.T) { + docker.ContainerStop(mdsName1) + }) + + t.Run("start a new mds node", func(t *testing.T) { + mc2 = featurebase.NewContainer(mdsName2, map[string]string{ + "FEATUREBASE_MDS_RUN": "true", + "FEATUREBASE_MDS_CONFIG_REGISTRATION_BATCH_TIMEOUT": "1ms", + "FEATUREBASE_STORAGE_METHOD": "boltdb", + "FEATUREBASE_STORAGE_DSN": dsn, + }) + _, err := docker.ContainerCreate(featurebase.ImageName, mc2, []docker.Volume{mdsVolume, coverVolume}) + assert.NoError(t, err) + assert.NoError(t, docker.NetworkConnect(networkName, mdsName2)) + assert.NoError(t, docker.ContainerStart(mdsName2)) + + waitForHealthy(t, mc2, dax.ServicePrefixMDS, 10, time.Second) + }) + // This closes the container created in "start a new compute node", but + // this should really only be called if that sub-test runs successfully. + defer func() { + docker.ContainerStop(mdsName2) + docker.ContainerRemove(mdsName2) + }() + + // ensure partitions are covered by the new mds node. + t.Run("check the translate-nodes endpoint again", func(t *testing.T) { + data := mdshttp.TranslateNodesRequest{ + Table: qtbl.QualifiedID(), + Partitions: partitions, + IsWrite: true, + } + + inspect(t, + mc2, + mdsAddress2, + "mds/translate-nodes", + data, + getExpFTranslateResponse(mdshttp.TranslateNodesResponse{ + TranslateNodes: []controller.TranslateNode{ + { + Address: addressFn(computerName1), + Partitions: dax.PartitionNums{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}, + }, + }}), + ) + }) + + // ensure the poller is polling the compute node(s) after the MDS restart. + + // stop featurebase 1 (may need to sleep) + t.Run("stop a container", func(t *testing.T) { + docker.ContainerStop(computerName1) + }) + + // Give the poller time to recognize the node is gone. + // TODO: implement this without a sleep. + time.Sleep(20 * time.Second) + + t.Run("start a new compute node", func(t *testing.T) { + fc2 := featurebase.NewContainer(computerName2, map[string]string{ + "FEATUREBASE_COMPUTER_RUN": "true", + "FEATUREBASE_COMPUTER_CONFIG_MDS_ADDRESS": string(mdsAddress2), + }) + _, err := docker.ContainerCreate(featurebase.ImageName, fc2, []docker.Volume{coverVolume}) + assert.NoError(t, err) + assert.NoError(t, docker.NetworkConnect(networkName, computerName2)) + assert.NoError(t, docker.ContainerStart(computerName2)) + + waitForHealthy(t, fc2, dax.ServicePrefixComputer, 10, time.Second) + waitForStatus(t, fc2, "NORMAL", 10, time.Second) + }) + // This closes the container created in "start a new compute node", but + // this should really only be called if that sub-test runs successfully. + defer func() { + docker.ContainerStop(computerName2) + docker.ContainerRemove(computerName2) + }() + + // ensure partitions are covered by the new compute node. + t.Run("check the translate-nodes endpoint after compute node shutdown", func(t *testing.T) { + data := mdshttp.TranslateNodesRequest{ + Table: qtbl.QualifiedID(), + Partitions: partitions, + IsWrite: true, + } + + inspect(t, + mc2, + mdsAddress2, + "mds/translate-nodes", + data, + getExpFTranslateResponse(mdshttp.TranslateNodesResponse{ + TranslateNodes: []controller.TranslateNode{ + { + Address: addressFn(computerName2), + Partitions: dax.PartitionNums{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}, + }, + }}), + ) + }) + }) + + // The idea for this test is to see if a node which has been paused (and + // therefore removed from the MDS's node list by the poller) will be + // re-registered via a CheckIn once it resumes following the pause. If this + // happens, querying the data should result in 0 records because we have not + // implemented a writelogger. This means the node re-joined, but it did not + // recover data from the writelogger, and what data it did have locally was + // removed due to a Directive.Method=reset, which is what we expect. + t.Run("Node_CheckIn", func(t *testing.T) { + imagePull(t, featurebase.ImageName) + imagePull(t, datagen.ImageName) + + tableName := dax.TableName("tbl") + partitionN := 256 + tbl := dax.NewTable(tableName) + tbl.PartitionN = partitionN + tbl.Fields = []*dax.Field{ + {Name: "_id", Type: "id"}, + {Name: "an_int", Type: "int", + Options: dax.FieldOptions{ + Min: pql.NewDecimal(0, 0), + Max: pql.NewDecimal(500, 0), + }, + }, + {Name: "a_random_string", Type: "string"}, + {Name: "an_id_set", Type: "idset"}, + {Name: "a_string_set", Type: "stringset"}, + } + qtbl := dax.NewQualifiedTable(qual, tbl) + + mcName := uniqueContainerName(t, mdsStub) + mc := featurebase.NewContainer(mcName, map[string]string{ + "FEATUREBASE_MDS_RUN": "true", + "FEATUREBASE_MDS_CONFIG_REGISTRATION_BATCH_TIMEOUT": "1ms", + }) + + fcName := uniqueContainerName(t, computerStub) + fc := featurebase.NewContainer(fcName, map[string]string{ + "FEATUREBASE_COMPUTER_RUN": "true", + "FEATUREBASE_COMPUTER_CONFIG_MDS_ADDRESS": string(addressFn(mcName)), + "FEATUREBASE_COMPUTER_CONFIG_CHECK_IN_INTERVAL": "2s", + }) + + qcName := uniqueContainerName(t, queryerStub) + qc := featurebase.NewContainer(qcName, map[string]string{ + "FEATUREBASE_QUERYER_RUN": "true", + "FEATUREBASE_QUERYER_CONFIG_MDS_ADDRESS": string(addressFn(mcName)), + }) + + networkName := uniqueNetworkName(t, mdsNetworkName) + + dc := new(docker.Composer). + WithVolume(coverVolume, fc, mc, qc). + WithService(featurebase.ImageName, mc). + WithService(featurebase.ImageName, fc). + WithService(featurebase.ImageName, qc). + WithNetwork(networkName) + + t.Run("start the containers with no errors", func(t *testing.T) { + require.NoError(dc.Up()) + waitForHealthy(t, mc, dax.ServicePrefixMDS, 10, time.Second) + waitForHealthy(t, qc, dax.ServicePrefixQueryer, 10, time.Second) + waitForHealthy(t, fc, dax.ServicePrefixComputer, 10, time.Second) + waitForStatus(t, fc, "NORMAL", 10, time.Second) + }) + defer dc.Down() + + t.Run("add a table", func(t *testing.T) { + sqlCheck(t, + mc, + addressFn(qcName), + qtbl.Qualifier(), + qtbl.CreateSQL(), + `{"schema":{"fields":[]},"data":[],"error":"","warnings":null}`, + ) + }) + + t.Run("ingest some data", func(t *testing.T) { + // datagen + assert.NoError(t, docker.Setenv("GEN_FEATUREBASE_ORG_ID", string(qual.OrganizationID))) + assert.NoError(t, docker.Setenv("GEN_FEATUREBASE_DB_ID", string(qual.DatabaseID))) + assert.NoError(t, docker.Setenv("GEN_FEATUREBASE_TABLE_NAME", string(qtbl.Name))) + assert.NoError(t, docker.Setenv("GEN_CUSTOM_CONFIG", "/testdata/basic.yaml")) + assert.NoError(t, docker.Setenv("GEN_USE_SHARD_TRANSACTIONAL_ENDPOINT", "true")) + gcName := uniqueContainerName(t, datagenStub) + gc := datagen.NewContainer(gcName, addressFn(mcName)) + _, err := docker.ContainerCreate(datagen.ImageName, gc, []docker.Volume{}) + assert.NoError(t, err) + assert.NoError(t, docker.NetworkConnect(networkName, gcName)) + assert.NoError(t, docker.ContainerStartWithLogging(gcName)) + defer func() { + docker.ContainerStop(gcName) + docker.ContainerRemove(gcName) + }() + + // Wait for datagen to finish ingest. + require.NoError(docker.ContainerWait(gc.Hostname())) + }) + + t.Run("query some data", func(t *testing.T) { + // PQL: `Count(Row(a_string_set="A90B"))`, + sqlCheck(t, + mc, + addressFn(qcName), + qual, + `select count(*) as cnt from tbl`, + `{"schema":{"fields":[{"name":"cnt","type":"INT"}]},"data":[[100]],"error":"","warnings":null}`, + ) + }) + + // stop featurebase 1 (may need to sleep) + t.Run("pause a container", func(t *testing.T) { + docker.ContainerPauseAndResume(fcName, 25*time.Second) + }) + + // Wait the 25 seconds for the pause/resume, plus another 25s. + // TODO: implement this without a sleep. + time.Sleep(50 * time.Second) + + t.Run("query some data again", func(t *testing.T) { + // PQL: `Count(Row(a_string_set="A90B"))`, + sqlCheck(t, + mc, + addressFn(qcName), + qual, + `select count(*) as cnt from tbl`, + `{"schema":{"fields":[{"name":"cnt","type":"INT"}]},"data":[[0]],"error":"","warnings":null}`, + ) + }) + }) +} + +type expFunc func(t *testing.T, actual string) + +// checks if the response contains all passed strings. +func getExpFContains(exps ...string) expFunc { + return func(t *testing.T, actual string) { + for _, exp := range exps { + assert.Contains(t, actual, exp) + } + } +} + +// compares if strings are qual +func getExpFEqual(exp string) expFunc { + return func(t *testing.T, actual string) { + assert.Equal(t, exp, actual) + } +} + +// compares tables, but only checks for presence of absence of ID, not +// actual value. +func getExpFTable(tabl *dax.QualifiedTable) expFunc { + return func(t *testing.T, actual string) { + actTabl := dax.Table{} + err := json.Unmarshal([]byte(actual), &actTabl) + assert.NoError(t, err) + assert.Equal(t, tabl.Name, actTabl.Name) + assert.Equal(t, tabl.Fields, actTabl.Fields) + assert.Equal(t, tabl.PartitionN, actTabl.PartitionN) + if len(tabl.ID) == 0 && len(actTabl.ID) != 0 { + t.Errorf("expected no ID, but got '%s'", actTabl.ID) + } + if len(tabl.ID) != 0 && len(actTabl.ID) == 0 { + t.Error("expected ID, but got none") + } + + for i, fld := range actTabl.Fields { + assert.Equal(t, tabl.Fields[i], fld) + } + } +} + +// compares translate nodes response, but doesn't compare node names. +func getExpFTranslateResponse(exp mdshttp.TranslateNodesResponse) expFunc { + return func(t *testing.T, actual string) { + actResp := mdshttp.TranslateNodesResponse{} + err := json.Unmarshal([]byte(actual), &actResp) + assert.NoError(t, err) + assert.Equal(t, len(exp.TranslateNodes), len(actResp.TranslateNodes)) + for i, actNode := range actResp.TranslateNodes { + expNode := exp.TranslateNodes[i] + assert.Equal(t, expNode.Address, actNode.Address) + assert.Equal(t, expNode.Partitions, actNode.Partitions) + } + } +} + +// inspect is a test helper for issuing common curl commands to a container and +// verifying the results are expected. +// +// c - the container on which the curl command should be run +// address - the target of the curl command +// data - if non-nil, will be marshaled to json as the POST arguments +// - if nil, the curl command will run as GET +// +// expF - A function called on the result which should return true if the result is as expected. +func inspect(t *testing.T, c docker.Container, address dax.Address, path string, data interface{}, expF expFunc) { + t.Helper() + + insp, err := inspector.NewInspector() + assert.NoError(t, err) + defer insp.Close() + + ctx := context.Background() + + req := inspector.NewExecRequest(address, path, data) + + cmd := req.Cmd() + log.Printf("cmd: %s", cmd) + + //resp, err := insp.ExecResp(ctx, c, req.Cmd()) + resp, err := insp.ExecResp(ctx, c, cmd) + assert.NoError(t, err) + expF(t, resp.Out()) + //assert.JSONEq(t, exp, resp.Out()) +} + +// sqlCheck is a test helper for issuing a sql command to a queryer and +// verifying the results are expected. +// +// c - the container on which the curl command should be run +// address - the queryer address +// qual - table qualifier +// sql - sql string +// +// exp - the expected result marshaled to a string +func sqlCheck(t *testing.T, c docker.Container, address dax.Address, qual dax.TableQualifier, sql string, exp string) { + t.Helper() + + out := sqlRun(t, c, address, qual, sql) + + got := &fb.SQLResponse{} + assert.NoError(t, json.Unmarshal([]byte(out), got)) + + want := &fb.SQLResponse{} + assert.NoError(t, json.Unmarshal([]byte(exp), want)) + + assert.Equal(t, want.Schema, got.Schema) + assert.Equal(t, want.Data, got.Data) + assert.Equal(t, want.Error, got.Error) + assert.Equal(t, want.Warnings, got.Warnings) +} + +// sqlRun is a test helper for issuing sql to a queryer's sql endpoint. +// +// c - the container on which the curl command should be run +// address - the queryer address +// qual - table qualifier +// sql - sql string +func sqlRun(tb testing.TB, c docker.Container, address dax.Address, qual dax.TableQualifier, sql string) string { + tb.Helper() + + path := "queryer/sql" + + insp, err := inspector.NewInspector() + assert.NoError(tb, err) + defer insp.Close() + + ctx := context.Background() + + sqlReq := queryerhttp.SQLRequest{ + OrganizationID: qual.OrganizationID, + DatabaseID: qual.DatabaseID, + SQL: sql, + } + + req := inspector.NewExecRequest(address, path, sqlReq) + + cmd := req.Cmd() + log.Printf("cmd: %s", cmd) + + //resp, err := insp.ExecResp(ctx, c, req.Cmd()) + resp, err := insp.ExecResp(ctx, c, cmd) + assert.NoError(tb, err) + + out := resp.Out() + log.Printf("RESPONSE: cmd: %s; out: %s", cmd, out) + + return out +} + +// pqlRun is a test helper for issuing pql to a queryer's sql endpoint. +// +// c - the container on which the curl command should be run +// address - the queryer address +// qual - table qualifier +// pql - pql string +func pqlRun(tb testing.TB, c docker.Container, address dax.Address, qual dax.TableQualifier, table, pql string) string { + tb.Helper() + + path := "queryer/query" + + insp, err := inspector.NewInspector() + assert.NoError(tb, err) + defer insp.Close() + + ctx := context.Background() + + pqlReq := queryerhttp.QueryRequest{ + OrganizationID: qual.OrganizationID, + DatabaseID: qual.DatabaseID, + Table: dax.TableName(table), + PQL: pql, + } + + req := inspector.NewExecRequest(address, path, pqlReq) + + cmd := req.Cmd() + log.Printf("cmd: %s", cmd) + + //resp, err := insp.ExecResp(ctx, c, req.Cmd()) + resp, err := insp.ExecResp(ctx, c, cmd) + assert.NoError(tb, err) + + out := resp.Out() + log.Printf("RESPONSE: cmd: %s; out: %s", cmd, out) + + return out +} + +// statusResponse mirrors the basic elements of getStatusResponse, a private +// type in the featurebase package. In the future, we should have shared, public +// API return types for featurebase. +type statusResponse struct { + State string `json:"state"` + LocalID string `json:"localID"` + ClusterName string `json:"clusterName"` +} + +// waitForStatus is currently specific to featurebase. In the future, we could +// generalize this by taking a docker.Container instead, and having every +// container support something like host:80/status (where we standardize on the +// port, /status, and the return payload) +func waitForStatus(t *testing.T, fc *featurebase.Container, status string, n int, sleep time.Duration) { + t.Helper() + + insp, err := inspector.NewInspector() + assert.NoError(t, err) + defer insp.Close() + + ctx := context.Background() + address := dax.Address(fc.Hostname() + ":8080") + req := inspector.NewExecRequest(address, fmt.Sprintf("%s/status", dax.ServicePrefixComputer), nil) + + for i := 0; i < n; i++ { + resp, err := insp.ExecResp(ctx, fc, req.Cmd()) + assert.NoError(t, err) + + var statusResp statusResponse + + if err := json.Unmarshal([]byte(resp.Out()), &statusResp); err != nil { + // treat json unmarshal error the same as not getting the expected + // status. + } else { + s := statusResp.State + t.Logf("Status (%d/%d): %s (sleep: %s)\n", i, n, s, sleep.String()) + + if s == status { + return + } + } + + if i < n-1 { + time.Sleep(sleep) + } + } + + // Getting to here means status was never found, so we need to stop the + // test. + panic("waitForStatus timed out") +} + +// waitForHealthy currently requires the service to be listening on port 8080. +// It polls the /health endpoint expecting `HTTP/1.1 200 OK`. +func waitForHealthy(t *testing.T, container docker.Container, path string, n int, sleep time.Duration) { + t.Helper() + + insp, err := inspector.NewInspector() + assert.NoError(t, err) + defer insp.Close() + + ctx := context.Background() + address := container.Hostname() + ":8080" + fullPath := "health" + if path != "" { + fullPath = path + "/" + fullPath + } + uri := fmt.Sprintf("%s/%s", address, fullPath) + cmd := []string{"curl", "-I", "-XGET", uri} + healthOk := "HTTP/1.1 200 OK" + + for i := 0; i < n; i++ { + resp, err := insp.ExecResp(ctx, container, cmd) + assert.NoError(t, err) + + out := resp.Out() + t.Logf("Health (%d/%d): uri: %s '%s' (sleep: %s)\n", i, n, uri, out, sleep.String()) + + if strings.HasPrefix(out, healthOk) { + return + } + + if i < n-1 { + time.Sleep(sleep) + } + } + + // Getting to here means the health endpoint never returned successfully, so + // we need to stop the test. + panic("waitForHealthy timed out") +} + +func uniqueNetworkName(t *testing.T, stub string) string { + rn := make([]byte, 8) + if _, err := rand.Read(rn); err != nil { + t.Fatalf("getting random data: %v", err) + } + return fmt.Sprintf("%s_%x", stub, rn) +} + +func uniqueContainerName(t *testing.T, stub string) string { + rn := make([]byte, 8) + if _, err := rand.Read(rn); err != nil { + t.Fatalf("getting random data: %v", err) + } + return fmt.Sprintf("%s%x", stub, rn) +} + +// sortStringKeys goes through an entire set of rows, and for any []string it +// finds, it orders the elements. This is obviously only useful in tests, and +// only in cases where we expect the elements to match, but we don't care what +// order they're in. It's basically the equivalent of assert.ElementsMatch(), +// but the way we use that on rows doesn't recurse down into the field values +// within each row. +// TODO(tlt): put this in sql test? +func sortStringKeys(in [][]interface{}) { + for i := range in { + for j := range in[i] { + switch v := in[i][j].(type) { + case []string: + sort.Strings(v) + } + } + } +} + +// mustQueryRows returns the row results as a slice of []interface{}, along with the columns. +func mustQueryRows(tb testing.TB, c docker.Container, address dax.Address, qual dax.TableQualifier, query string, table string) ([][]interface{}, []*planner_types.PlannerColumn, error) { + tb.Helper() + var sql string + var out string + if table == "" { + sql = query + out = sqlRun(tb, c, address, qual, sql) + } else { + out = pqlRun(tb, c, address, qual, table, query) + } + + sqlResp := &fb.SQLResponse{} + + err := json.Unmarshal([]byte(out), sqlResp) + if err != nil { + tb.Fatalf("error unmarshaling response: %v, raw resp: '%s'", err, out) + } + + headers := make([]*planner_types.PlannerColumn, len(sqlResp.Schema.Fields)) + for i, fld := range sqlResp.Schema.Fields { + var htype parser.ExprDataType + + if strings.HasPrefix(fld.Type, parser.FieldTypeDecimal) { + sscale := fld.Type[8 : len(fld.Type)-1] + scale, err := strconv.Atoi(sscale) + assert.NoError(tb, err) + htype = &parser.DataTypeDecimal{ + Scale: int64(scale), + } + + } else { + switch fld.Type { + case parser.FieldTypeBool: + htype = &parser.DataTypeBool{} + case parser.FieldTypeID: + htype = &parser.DataTypeID{} + case parser.FieldTypeIDSet: + htype = &parser.DataTypeIDSet{} + case parser.FieldTypeIDSetQuantum: + htype = &parser.DataTypeIDSetQuantum{} + case parser.FieldTypeInt: + htype = &parser.DataTypeInt{} + case parser.FieldTypeString: + htype = &parser.DataTypeString{} + case parser.FieldTypeStringSet: + htype = &parser.DataTypeStringSet{} + case parser.FieldTypeStringSetQuantum: + htype = &parser.DataTypeStringSetQuantum{} + case parser.FieldTypeTimestamp: + htype = &parser.DataTypeTimestamp{} + default: + tb.Errorf("unsupported header type: %s", fld.Type) + } + } + + headers[i] = &planner_types.PlannerColumn{ + ColumnName: fld.Name, + Type: htype, + } + } + + data := sqlResp.Data + + // try to convert the types based on the headers + for i := range data { + for j, hdr := range headers { + switch ht := hdr.Type.(type) { + case *parser.DataTypeID, *parser.DataTypeInt: + if _, ok := data[i][j].(float64); ok { + data[i][j] = int64(data[i][j].(float64)) + } + + case *parser.DataTypeIDSet: + if src, ok := data[i][j].([]interface{}); ok { + val := make([]int64, len(src)) + for k := range src { + val[k] = int64(src[k].(float64)) + } + data[i][j] = val + } + + case *parser.DataTypeDecimal: + if _, ok := data[i][j].(float64); ok { + format := fmt.Sprintf("%%.%df", ht.Scale) + dec, err := pql.ParseDecimal(fmt.Sprintf(format, data[i][j])) + assert.NoError(tb, err) + data[i][j] = dec + } + + case *parser.DataTypeStringSet: + if src, ok := data[i][j].([]interface{}); ok { + val := make([]string, len(src)) + for k := range src { + val[k] = src[k].(string) + } + data[i][j] = val + } + + case *parser.DataTypeBool, *parser.DataTypeString: + // no need to convert + + default: + log.Printf("WARNING: unimplemented: %T", ht) + } + } + } + + if sqlResp.Error != "" { + err = errors.New(sqlResp.Error) + } + + return data, headers, err +} diff --git a/dax/test/dax/image.go b/dax/test/dax/image.go new file mode 100644 index 000000000..70b958099 --- /dev/null +++ b/dax/test/dax/image.go @@ -0,0 +1,24 @@ +package dax + +import ( + "strings" + "testing" + + "github.com/molecula/featurebase/v3/dax/test/docker" + "github.com/stretchr/testify/require" +) + +func imagePull(t *testing.T, imageName ...string) { + t.Helper() + for _, img := range imageName { + // If the images is something other than one at docker.io, then don't + // perform the `docker pull`. This allows a develper to swap out the + // ImageName with a local image. + if !strings.HasPrefix(img, "docker.io/") { + continue + } + if err := docker.ImagePull(img); err != nil { + require.NoError(t, err) + } + } +} diff --git a/dax/test/dax/testdata/basic.yaml b/dax/test/dax/testdata/basic.yaml new file mode 100644 index 000000000..aca6bfaa5 --- /dev/null +++ b/dax/test/dax/testdata/basic.yaml @@ -0,0 +1,52 @@ +fields: + - name: "an_id" + type: "uint" # (default IDField (non-mutex)) + distribution: "sequential" + min: 0 + max: 2000000 + repeat: false # if false, data generation stops when we hit >= max. only available with sequential + step: 1 + - name: "an_int" + type: "int" # (default IntField) + distribution: "uniform" # uniform or zipfian + min: 0 + max: 500 + - name: "a_random_string" + type: "string" # (default StringField (non-mutex)) + generator_type: "random-string" # used to generate random strings rather than pulling from known set + min_len: 3 + max_len: 3 + charset: "AB" # set of possible characters to pull from when generating random string + - name: "an_id_set" + type: "uint-set" # (default IDArrayField) + min: 0 + max: 1000 + distribution: "uniform" + min_num: 1 + max_num: 6 + - name: "a_string_set" + type: "string-set" # (default StringArrayField) + generator_type: "random-string" # used to generate random strings rather than pulling from known. "distribution" is ignored. + min_len: 4 + max_len: 4 + charset: "0123456789ABCDEF" + min_num: 0 # minimum number of strings in each value (default 0) + max_num: 10 # max number of strings (default to cardinality of source) + +# idk_params describe how data from "fields" should be ingested by IDK +idk_params: + primary_key_config: + field: "id" # if this is a single field named "id" then we'll use uint IDs, if it's empty we'll autogen ids, and if it's anything else we'll do string keys... yes this is a bit hacky, needs to be cleaned up. + # fields is keyed by names of fields from top level "fields". It is + # not required that all fields appear here, those that don't will + # use the default ingestion. + fields: + an_id: + - type: "ID" + mutex: false + name: "id" + a_string_set: + - type: "StringArray" + a_random_string: + - type: "String" + mutex: true diff --git a/dax/test/docker/client.go b/dax/test/docker/client.go new file mode 100644 index 000000000..1d2f6a4d7 --- /dev/null +++ b/dax/test/docker/client.go @@ -0,0 +1,34 @@ +package docker + +import ( + "context" + "os" + + "github.com/docker/docker/api/types" + "github.com/docker/docker/client" +) + +// DefaultClient is the default docker Client +var DefaultClient *client.Client + +func init() { + cli, err := client.NewClientWithOpts() + if err != nil { + panic(err) + } + cli.RegistryLogin(context.Background(), types.AuthConfig{}) + + DefaultClient = cli +} + +func Getenv(key, fallback string) string { + value := os.Getenv(key) + if len(value) == 0 { + return fallback + } + return value +} + +func Setenv(key, value string) error { + return os.Setenv(key, value) +} diff --git a/dax/test/docker/docker.go b/dax/test/docker/docker.go new file mode 100644 index 000000000..518e92321 --- /dev/null +++ b/dax/test/docker/docker.go @@ -0,0 +1,521 @@ +package docker + +import ( + "bufio" + "context" + "fmt" + "io" + "log" + "os/exec" + "strings" + "syscall" + "time" + + "github.com/docker/docker/api/types" + "github.com/docker/docker/api/types/container" + "github.com/docker/docker/api/types/filters" + "github.com/docker/docker/api/types/mount" + "github.com/docker/docker/api/types/network" + "github.com/docker/docker/api/types/strslice" + "github.com/docker/docker/api/types/volume" + "github.com/docker/docker/pkg/stdcopy" + "github.com/docker/go-connections/nat" + "github.com/molecula/featurebase/v3/errors" + v1 "github.com/opencontainers/image-spec/specs-go/v1" +) + +type Composer struct { + network string + containers []Container + + // volumes is a map of container to volume(s). + volumes map[string][]Volume +} + +type Volume struct { + Type string + Source string + Target string +} + +// WithVolume creates a volume and registers it with the provided containers. +// Note that WithVolume must be called before WithService for the applicable +// containers. +func (c *Composer) WithVolume(volume Volume, containers ...Container) *Composer { + if c.volumes == nil { + c.volumes = make(map[string][]Volume) + } + + // delete previous existing networks if any to avoid creation errors + vol, _ := VolumeByName(volume.Source) + if vol != nil { + VolumeRemove(vol.Name) + VolumePrune(vol.Name) + } + + // TODO(jaffee) I don't understand what the other volume type is + // being used for, but if you try to pass a bind mount here it + // panics, so I put a janky "if" around it. + if volume.Type != "bind" { + _, err := VolumeCreate(volume.Source) + if err != nil { + panic(err) + } + } + + // Register volume with container(s). + for _, cont := range containers { + c.volumes[cont.Hostname()] = append(c.volumes[cont.Hostname()], volume) + } + + return c +} + +func (c *Composer) WithNetwork(networkName string) *Composer { + // delete previous existing networks if any to avoid creation errors + net, _ := NetworkByName(networkName) + NetworkRemove(net.Name) + NetworkPrune(net.Name) + + _, err := NetworkCreate(networkName) + if err != nil { + panic(err) + } + + for _, c := range c.containers { + if err = NetworkConnect(networkName, c.Hostname()); err != nil { + panic(err) + } + } + c.network = networkName + + return c +} + +func (c *Composer) WithService(imageName string, container ...Container) *Composer { + for _, cc := range container { + // remove previous containers with the same name + pc, _ := ContainerByName(cc.Hostname()) + for _, n := range pc.Names { + ContainerStop(n) + ContainerRemove(n) + ContainerPrune(n) + } + + _, err := ContainerCreate(imageName, cc, c.volumes[cc.Hostname()]) + if err != nil { + panic(err) + } + c.containers = append(c.containers, cc) + if c.network != "" { + if err = NetworkConnect(c.network, cc.Hostname()); err != nil { + panic(err) + } + } + } + + return c +} + +func (c *Composer) Up() error { + for _, cc := range c.containers { + if err := ContainerStartWithLogging(cc.Hostname()); err != nil { + return err + } + } + + return nil +} + +func (c *Composer) Down() error { + var errs []error + + for _, cc := range c.containers { + name := cc.Hostname() + + if err := ContainerStop(name); err != nil { + log.Printf("Composer.Down error: ContainerStop: %s: %v", name, err) + errs = append(errs, err) + } + + if err := ContainerRemove(name); err != nil { + log.Printf("Composer.Down error: ContainerRemove: %s: %v", name, err) + errs = append(errs, err) + } + } + + if c.network != "" { + if err := NetworkRemove(c.network); err != nil { + log.Printf("Composer.Down error: NetworkRemove: %v", err) + errs = append(errs, err) + } + } + + if len(errs) == 0 { + return nil + } + + var errString strings.Builder + + for i := range errs { + errString.WriteString(fmt.Sprintf("(%d) ", i)) + errString.WriteString(errs[i].Error()) + errString.WriteString(" ") + } + + return errors.New(errors.ErrUncoded, errString.String()) +} + +type Container interface { + Hostname() string + ExposedPorts() []string + Env() map[string]string + Cmd() []string +} + +// ImagePull pulls the selected image from internet +func ImagePull(imageName string) error { + // TODO temporal. Problems with authorization + cmd := exec.Command("docker", "pull", imageName) + if err := cmd.Run(); err != nil { + if exiterr, ok := err.(*exec.ExitError); ok { + if status, ok := exiterr.Sys().(syscall.WaitStatus); ok && status.ExitStatus() > 0 { + return err + } + } else { + return err + } + } + return nil +} + +func ContainerCreate(imageName string, c Container, volumes []Volume) (string, error) { + exposedPorts := make(nat.PortSet) + for _, p := range c.ExposedPorts() { + exposedPorts[nat.Port(p)] = struct{}{} + } + + var platform *v1.Platform + + mnts := make([]mount.Mount, 0) + for _, v := range volumes { + typ := v.Type + if typ == "" { + typ = string(mount.TypeVolume) + } + mnts = append(mnts, mount.Mount{ + Type: mount.Type(typ), + Source: v.Source, + Target: v.Target, + }) + } + + resp, err := DefaultClient.ContainerCreate(context.Background(), + &container.Config{ + Image: imageName, + Env: envToSlice(c.Env()), + Hostname: c.Hostname(), + ExposedPorts: exposedPorts, + Cmd: strslice.StrSlice(c.Cmd()), + }, + &container.HostConfig{ + // AutoRemove: true, + NetworkMode: "bridge", + PublishAllPorts: true, + Mounts: mnts, + }, + &network.NetworkingConfig{}, platform, c.Hostname()) + if err != nil { + return "", err + } + return resp.ID, err +} + +func ContainerByName(containerName string) (types.Container, error) { + f := filters.NewArgs() + f.Add("name", containerName) + + l, err := DefaultClient.ContainerList(context.Background(), types.ContainerListOptions{ + Limit: 1, + Filters: f, + }) + if err != nil { + return types.Container{}, err + } + if len(l) == 0 { + return types.Container{}, fmt.Errorf("Container %s not found", containerName) + } + + return l[0], nil +} + +func ContainerPrune(containerName string) error { + f := filters.NewArgs() + f.Add("name", containerName) + + _, err := DefaultClient.ContainersPrune(context.Background(), f) + return err +} + +func ContainerStart(containerName string) error { + c, err := ContainerByName(containerName) + if err != nil { + return err + } + return DefaultClient.ContainerStart(context.Background(), c.ID, types.ContainerStartOptions{}) +} + +func ContainerStartWithLogging(containerName string) error { + if err := ContainerStart(containerName); err != nil { + return err + } + + go ContainerLogs(containerName) + + return nil +} + +func ContainerWait(containerName string) error { + c, err := ContainerByName(containerName) + if err != nil { + return err + } + + var cond container.WaitCondition = container.WaitConditionNotRunning + + statusCh, errCh := DefaultClient.ContainerWait(context.Background(), c.ID, cond) + _ = errCh + + status := <-statusCh + + log.Printf("ContainerWait status code: %d", status.StatusCode) + if status.Error != nil { + log.Printf("ContainerWait error: %s", status.Error.Message) + } + + if err != nil { + return errors.WithMessagef(err, "ContainerWait(%s: %s) error status code: %v", c.ID, containerName, status) + } + return nil +} + +func ContainerRestart(containerName string) error { + c, err := ContainerByName(containerName) + if err != nil { + return err + } + + return DefaultClient.ContainerRestart(context.Background(), c.ID, nil) +} + +func ContainerPauseAndResume(containerName string, timeout time.Duration) (err error) { + c, err := ContainerByName(containerName) + if err != nil { + return err + } + + if err = DefaultClient.ContainerPause(context.Background(), c.ID); err != nil { + return errors.WithMessagef(err, "ContainerPauseAndResume(%s: %s) error", c.ID, containerName) + } + + time.AfterFunc(timeout, func() { + if err = DefaultClient.ContainerUnpause(context.Background(), c.ID); err != nil { + err = errors.WithMessagef(err, "ContainerUnpause(%s: %s) error", c.ID, containerName) + } + }) + + return err +} + +func ContainerLogs(containerName string) error { + c, err := ContainerByName(containerName) + if err != nil { + return err + } + + reader, err := DefaultClient.ContainerLogs(context.Background(), c.ID, + types.ContainerLogsOptions{ + ShowStderr: true, + ShowStdout: true, + Follow: true, + }) + if err != nil { + return err + } + defer reader.Close() + + r, w := io.Pipe() + + go func() { + stdcopy.StdCopy(w, w, reader) + }() + + br := bufio.NewReader(r) + for { + line, err := br.ReadString('\n') + + if err == io.EOF { + break + } + if err != nil { + return err + } + + fmt.Print("LOGS from ", containerName, ": ", line) + } + + return nil +} + +func ContainerStop(containerName string) error { + c, err := ContainerByName(containerName) + if err != nil { + return err + } + + return DefaultClient.ContainerStop(context.Background(), c.ID, nil) +} + +func ContainerRemove(containerName string) error { + c, err := ContainerByName(containerName) + if err != nil { + return err + } + + return DefaultClient.ContainerRemove( + context.Background(), + c.ID, + types.ContainerRemoveOptions{ + Force: true, + RemoveVolumes: true, + }, + ) +} + +func VolumeCreate(volumeName string) (string, error) { + res, err := DefaultClient. + VolumeCreate( + context.Background(), + volume.VolumeCreateBody{ + Name: volumeName, + }, + ) + + return res.Name, err +} + +func VolumeByName(volumeName string) (*types.Volume, error) { + f := filters.NewArgs() + f.Add("name", volumeName) + + n, err := DefaultClient.VolumeList(context.Background(), f) + if err != nil { + return nil, err + } + if len(n.Volumes) == 0 { + return nil, fmt.Errorf("volume %s not found", volumeName) + } + + return n.Volumes[0], nil +} + +func VolumeRemove(volumeName string) error { + n, err := VolumeByName(volumeName) + if err != nil { + return err + } + + return DefaultClient.VolumeRemove(context.Background(), n.Name, false) +} + +func VolumePrune(volumeName string) error { + f := filters.NewArgs() + f.Add("name", volumeName) + + _, err := DefaultClient.VolumesPrune(context.Background(), f) + return err +} + +func NetworkCreate(networkName string) (string, error) { + res, err := DefaultClient. + NetworkCreate( + context.Background(), + networkName, + types.NetworkCreate{ + CheckDuplicate: true, + }, + ) + + return res.ID, err +} + +func NetworkByName(networkName string) (types.NetworkResource, error) { + f := filters.NewArgs() + f.Add("name", networkName) + + n, err := DefaultClient.NetworkList(context.Background(), types.NetworkListOptions{ + Filters: f, + }) + if err != nil { + return types.NetworkResource{}, err + } + if len(n) == 0 { + return types.NetworkResource{}, fmt.Errorf("network %s not found", networkName) + } + + return n[0], nil +} + +func NetworkRemove(networkName string) error { + n, err := NetworkByName(networkName) + if err != nil { + return err + } + + return DefaultClient.NetworkRemove(context.Background(), n.ID) +} + +func NetworkConnect(networkName, containerName string) error { + n, err := NetworkByName(networkName) + if err != nil { + return err + } + + c, err := ContainerByName(containerName) + if err != nil { + return err + } + + return DefaultClient.NetworkConnect(context.Background(), n.ID, c.ID, &network.EndpointSettings{}) +} + +func NetworkDisconnect(networkName, containerName string) error { + n, err := NetworkByName(networkName) + if err != nil { + return err + } + + c, err := ContainerByName(containerName) + if err != nil { + return err + } + + return DefaultClient.NetworkDisconnect(context.Background(), n.ID, c.ID, true) +} + +func NetworkPrune(networkName string) error { + f := filters.NewArgs() + f.Add("name", networkName) + + _, err := DefaultClient.NetworksPrune(context.Background(), f) + return err +} + +func envToSlice(env map[string]string) []string { + var out = make([]string, 0, len(env)) + for k, v := range env { + out = append(out, fmt.Sprintf("%s=%s", k, v)) + } + + return out +} diff --git a/dax/test/featurebase/container.go b/dax/test/featurebase/container.go new file mode 100644 index 000000000..921a9c443 --- /dev/null +++ b/dax/test/featurebase/container.go @@ -0,0 +1,75 @@ +package featurebase + +import ( + "fmt" + + "github.com/molecula/featurebase/v3/dax/test/docker" +) + +var ImageName = docker.Getenv("FEATUREBASE_DOCKER_IMAGE", "dax/featurebase-test") + +const ( + DataDir = "/data" + NetworkName = "featurebase-network" + + HTTPPort = "8080" + GRPCPort = "20101" + AdvertisePeerAddr = "2379" + AdvertiseClientAddr = "2380" +) + +// Ensure type implements interface. +var _ docker.Container = &Container{} + +type Container struct { + name string + replica int + cmd []string + env map[string]string +} + +func NewContainer(name string, env map[string]string) *Container { + // peers is just "self" because we don't want to use etcd as a cluster in + // the case of dumb compute nodes. + // peers := fmt.Sprintf("%s=http://%s:%s", name, name, AdvertisePeerAddr) + + c := &Container{ + name: name, + replica: 1, + cmd: []string{ + "/featurebase", + "-test.run=TestRunMain", + fmt.Sprintf("-test.coverprofile=/results/coverage-%s.out", name), + "dax", + }, + env: map[string]string{ + "FEATUREBASE_BIND": "0.0.0.0:" + HTTPPort, + "FEATUREBASE_ADVERTISE": name + ":" + HTTPPort, + }, + } + + // Apply given env vars. + for k, v := range env { + c.env[k] = v + } + + return c +} + +func (c *Container) Hostname() string { + return c.name +} + +func (c *Container) ExposedPorts() []string { + // We don't expose HTTPPort, because it's already exposed by default by + // featurebase docker image. + return []string{GRPCPort, AdvertisePeerAddr, AdvertiseClientAddr} +} + +func (c *Container) Cmd() []string { + return c.cmd +} + +func (c *Container) Env() map[string]string { + return c.env +} diff --git a/dax/test/inspector/inspector.go b/dax/test/inspector/inspector.go new file mode 100644 index 000000000..53141dec6 --- /dev/null +++ b/dax/test/inspector/inspector.go @@ -0,0 +1,217 @@ +package inspector + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "log" + "strings" + + "github.com/docker/docker/api/types" + "github.com/docker/docker/client" + "github.com/docker/docker/pkg/stdcopy" + "github.com/molecula/featurebase/v3/dax" + mdshttp "github.com/molecula/featurebase/v3/dax/mds/http" + queryerhttp "github.com/molecula/featurebase/v3/dax/queryer/http" + "github.com/molecula/featurebase/v3/dax/test/docker" + "github.com/molecula/featurebase/v3/errors" +) + +type Inspector struct { + cli *client.Client + + // containers is a map of container name to container ID. + containers map[string]string +} + +func NewInspector() (*Inspector, error) { + cli, err := client.NewClientWithOpts(client.FromEnv) + if err != nil { + return nil, errors.Wrap(err, "getting new client with opts") + } + + containers, err := cli.ContainerList(context.Background(), types.ContainerListOptions{}) + if err != nil { + return nil, errors.Wrap(err, "getting container list") + } + + m := make(map[string]string) + for _, container := range containers { + if len(container.Names) == 0 { + continue + } + m[container.Names[0]] = container.ID[:10] + } + + return &Inspector{ + cli: cli, + containers: m, + }, nil +} + +type ExecRequest struct { + address dax.Address + path string + data interface{} +} + +func NewExecRequest(address dax.Address, path string, data interface{}) *ExecRequest { + return &ExecRequest{ + address: address, + path: path, + data: data, + } +} + +func (e *ExecRequest) Cmd() []string { + address := e.address + if address == "" { + address = "http://localhost:8080" + } + uri := fmt.Sprintf("%s/%s", address, e.path) + + if e.data == nil { + return []string{ + "curl", + uri, + } + } + + var dArgs string + switch v := e.data.(type) { + case string: + dArgs = v + case *dax.Table, + *dax.QualifiedTable, + mdshttp.RegisterNodeRequest, + mdshttp.SnapshotFieldKeysRequest, + mdshttp.SnapshotShardRequest, + mdshttp.TranslateNodesRequest, + queryerhttp.QueryRequest, + queryerhttp.SQLRequest, + dax.QualifiedTableID: + b, err := json.Marshal(v) + if err != nil { + panic(err) // FIX THIS + } + dArgs = string(b) + default: + log.Printf("unhandled return type: %T", e.data) + dArgs = "{}" + } + + return []string{ + "curl", + "-d " + dArgs, + uri, + } +} + +type ExecResponse struct { + stdOut string + stdErr string + exitCode int +} + +func (e *ExecResponse) Out() string { + return strings.TrimSpace(e.stdOut) +} + +func (e *ExecResponse) Err() string { + return strings.TrimSpace(e.stdErr) +} + +func (i *Inspector) Close() error { + if i.cli != nil { + return i.cli.Close() + } + return nil +} + +// ExecResp is a helper method which runs exec() then resp(). +func (i *Inspector) ExecResp(ctx context.Context, container docker.Container, command []string) (ExecResponse, error) { + var execResp ExecResponse + + exec, err := i.exec(ctx, container, command) + if err != nil { + return execResp, err + } + + return i.resp(context.Background(), exec.ID) +} + +func (i *Inspector) exec(ctx context.Context, container docker.Container, command []string) (types.IDResponse, error) { + hostName := container.Hostname() + + containerID, ok := i.containers[slash(hostName)] + if !ok { + return types.IDResponse{}, errors.Errorf("invalid container: %s", hostName) + } + + config := types.ExecConfig{ + AttachStderr: true, + AttachStdout: true, + Cmd: command, + } + + return i.cli.ContainerExecCreate(ctx, containerID, config) +} + +func (i *Inspector) resp(ctx context.Context, id string) (ExecResponse, error) { + var execResp ExecResponse + + resp, err := i.cli.ContainerExecAttach(ctx, id, types.ExecStartCheck{}) + if err != nil { + return execResp, err + } + defer resp.Close() + + // read the output + var outBuf, errBuf bytes.Buffer + outputDone := make(chan error) + + go func() { + // StdCopy demultiplexes the stream into two buffers + _, err = stdcopy.StdCopy(&outBuf, &errBuf, resp.Reader) + outputDone <- err + }() + + select { + case err := <-outputDone: + if err != nil { + return execResp, err + } + break + + case <-ctx.Done(): + return execResp, ctx.Err() + } + + stdout, err := io.ReadAll(&outBuf) + if err != nil { + return execResp, err + } + stderr, err := io.ReadAll(&errBuf) + if err != nil { + return execResp, err + } + + res, err := i.cli.ContainerExecInspect(ctx, id) + if err != nil { + return execResp, err + } + + execResp.exitCode = res.ExitCode + execResp.stdOut = string(stdout) + execResp.stdErr = string(stderr) + return execResp, nil +} + +// slash is a helper functions which addresses the fact that the lower level +// docker inspect functions actually store the container name with a leading +// slash. See this for more info: https://github.com/moby/moby/issues/6705 +func slash(s string) string { + return "/" + s +} diff --git a/dax/test/schemar.go b/dax/test/schemar.go new file mode 100644 index 000000000..6002de722 --- /dev/null +++ b/dax/test/schemar.go @@ -0,0 +1,29 @@ +package test + +import ( + "os" + "testing" + + "github.com/molecula/featurebase/v3/dax/boltdb" + "github.com/molecula/featurebase/v3/dax/mds/schemar" + schemarbolt "github.com/molecula/featurebase/v3/dax/mds/schemar/boltdb" + testbolt "github.com/molecula/featurebase/v3/dax/test/boltdb" + "github.com/molecula/featurebase/v3/logger" +) + +func NewSchemar(t *testing.T) (schemar schemar.Schemar, cleanup func()) { + td, err := os.MkdirTemp("", "schemartest_*") + if err != nil { + t.Fatalf(": %v", err) + } + db, err := boltdb.NewSvcBolt(td, "schemar", schemarbolt.SchemarBuckets...) + if err != nil { + t.Fatalf("opening boltdb: %v", err) + } + + s := schemarbolt.NewSchemar(db, logger.StderrLogger) + return s, func() { + testbolt.MustCloseDB(t, db) + testbolt.CleanupDB(t, db.Path()) + } +} diff --git a/dax/test/table.go b/dax/test/table.go new file mode 100644 index 000000000..3c400c127 --- /dev/null +++ b/dax/test/table.go @@ -0,0 +1,69 @@ +// Package test include external test apps, helper functions, and test data. +package test + +import ( + "testing" + + "github.com/molecula/featurebase/v3/dax" +) + +// TestQualifiedTable is a test helper function for creating a table based on a +// general configuration. This function creates a Table with a random TableID. +// If you need to specify the TableID yourself, use the TestQualifiedTableWithID +// function. +func TestQualifiedTable(t *testing.T, qual dax.TableQualifier, name dax.TableName, partitionN int, keyed bool) *dax.QualifiedTable { + t.Helper() + + var pkFieldType dax.FieldType + if keyed { + pkFieldType = dax.FieldTypeString + } else { + pkFieldType = dax.FieldTypeID + } + + tbl := dax.NewTable(name) + tbl.CreateID() + tbl.PartitionN = partitionN + tbl.Fields = []*dax.Field{ + { + Name: dax.PrimaryKeyFieldName, + Type: pkFieldType, + }, + } + + return dax.NewQualifiedTable( + qual, + tbl, + ) +} + +// TestQualifiedTableWithID is a test helper function for creating a table based +// on a general configuration, and having the specified TableID. +func TestQualifiedTableWithID(t *testing.T, qual dax.TableQualifier, id string, name dax.TableName, partitionN int, keyed bool) *dax.QualifiedTable { + t.Helper() + + var pkFieldType dax.FieldType + if keyed { + pkFieldType = dax.FieldTypeString + } else { + pkFieldType = dax.FieldTypeID + } + + tbl := &dax.Table{ + ID: dax.TableID(id), + Name: name, + } + + tbl.PartitionN = partitionN + tbl.Fields = []*dax.Field{ + { + Name: dax.PrimaryKeyFieldName, + Type: pkFieldType, + }, + } + + return dax.NewQualifiedTable( + qual, + tbl, + ) +} diff --git a/dax/time.go b/dax/time.go new file mode 100644 index 000000000..d13e9ac0d --- /dev/null +++ b/dax/time.go @@ -0,0 +1,62 @@ +package dax + +import ( + "strings" +) + +// TimeQuantum represents a time granularity for time-based bitmaps. +type TimeQuantum string + +// HasYear returns true if the quantum contains a 'Y' unit. +func (q TimeQuantum) HasYear() bool { return strings.ContainsRune(string(q), 'Y') } + +// HasMonth returns true if the quantum contains a 'M' unit. +func (q TimeQuantum) HasMonth() bool { return strings.ContainsRune(string(q), 'M') } + +// HasDay returns true if the quantum contains a 'D' unit. +func (q TimeQuantum) HasDay() bool { return strings.ContainsRune(string(q), 'D') } + +// HasHour returns true if the quantum contains a 'H' unit. +func (q TimeQuantum) HasHour() bool { return strings.ContainsRune(string(q), 'H') } + +// IsEmpty returns true if the quantum is empty. +func (q TimeQuantum) IsEmpty() bool { return string(q) == "" } + +func (q TimeQuantum) Granularity() rune { + var g rune + for _, g = range q { + } + return g +} + +// Valid returns true if q is a valid time quantum value. +func (q TimeQuantum) Valid() bool { + switch q { + case "Y", "YM", "YMD", "YMDH", + "M", "MD", "MDH", + "D", "DH", + "H", + "": + return true + default: + return false + } +} + +// The following methods are required to implement pflag Value interface. + +// Set sets the time quantum value. +func (q *TimeQuantum) Set(value string) error { + *q = TimeQuantum(value) + return nil +} + +// String returns the TimeQuantum as a string type. +func (q TimeQuantum) String() string { + return string(q) +} + +// Type returns the type of a time quantum value. +func (q TimeQuantum) Type() string { + return "TimeQuantum" +} diff --git a/dax/versionstore.go b/dax/versionstore.go new file mode 100644 index 000000000..112cdd00e --- /dev/null +++ b/dax/versionstore.go @@ -0,0 +1,110 @@ +package dax + +import ( + "context" +) + +// VersionStore is an interface for tracking Shard, Partition, and Field[Key] +// versions. For example, when the contents of a shard are checkpointed, and a +// snapshot is generated, and the write log messages for that shard are +// truncated, the ShardVersion for that shard is incremented. The VersionStore +// is the interface through which various services read/write that version. +type VersionStore interface { + AddTable(ctx context.Context, qtid QualifiedTableID) error + RemoveTable(ctx context.Context, qtid QualifiedTableID) (Shards, Partitions, error) + + // Shards (shardData) + AddShards(ctx context.Context, qtid QualifiedTableID, shards ...Shard) error + Shards(ctx context.Context, qtid QualifiedTableID) (Shards, bool, error) + ShardVersion(ctx context.Context, qtid QualifiedTableID, shardNum ShardNum) (int, bool, error) + ShardTables(ctx context.Context, qual TableQualifier) (TableIDs, error) + + // Partitions (tableKeys) + AddPartitions(ctx context.Context, qtid QualifiedTableID, partitions ...Partition) error + Partitions(ctx context.Context, qtid QualifiedTableID) (Partitions, bool, error) + PartitionVersion(ctx context.Context, qtid QualifiedTableID, partitionNum PartitionNum) (int, bool, error) + PartitionTables(ctx context.Context, qual TableQualifier) (TableIDs, error) + + // Fields (fieldKeys) + AddFields(ctx context.Context, qtid QualifiedTableID, fields ...FieldVersion) error + Fields(ctx context.Context, qtid QualifiedTableID) (FieldVersions, bool, error) + FieldVersion(ctx context.Context, qtid QualifiedTableID, field FieldName) (int, bool, error) + FieldTables(ctx context.Context, qual TableQualifier) (TableIDs, error) + + Copy(ctx context.Context) (VersionStore, error) +} + +type DirectiveVersion interface { + Increment(ctx context.Context, delta uint64) (uint64, error) +} + +// Ensure type implements interface. +var _ VersionStore = (*nopVersionStore)(nil) + +// nopVersionStore is a no-op implementation of the VersionStore interface. +type nopVersionStore struct{} + +// NewNopVersionStore returns a new no-op instance of VersionStore. +func NewNopVersionStore() *nopVersionStore { + return &nopVersionStore{} +} + +func (s *nopVersionStore) AddTable(ctx context.Context, qtid QualifiedTableID) error { + return nil +} + +func (s *nopVersionStore) RemoveTable(ctx context.Context, qtid QualifiedTableID) (Shards, Partitions, error) { + return nil, nil, nil +} + +func (s *nopVersionStore) AddShards(ctx context.Context, qtid QualifiedTableID, shards ...Shard) error { + return nil +} + +func (s *nopVersionStore) Shards(ctx context.Context, qtid QualifiedTableID) (Shards, bool, error) { + return nil, false, nil +} + +func (s *nopVersionStore) ShardVersion(ctx context.Context, qtid QualifiedTableID, shardNum ShardNum) (int, bool, error) { + return 0, true, nil +} + +func (s *nopVersionStore) ShardTables(ctx context.Context, qual TableQualifier) (TableIDs, error) { + return TableIDs{}, nil +} + +func (s *nopVersionStore) AddPartitions(ctx context.Context, qtid QualifiedTableID, partitions ...Partition) error { + return nil +} + +func (s *nopVersionStore) Partitions(ctx context.Context, qtid QualifiedTableID) (Partitions, bool, error) { + return nil, false, nil +} + +func (s *nopVersionStore) PartitionVersion(ctx context.Context, qtid QualifiedTableID, partitionNum PartitionNum) (int, bool, error) { + return 0, true, nil +} + +func (s *nopVersionStore) PartitionTables(ctx context.Context, qual TableQualifier) (TableIDs, error) { + return TableIDs{}, nil +} + +func (s *nopVersionStore) AddFields(ctx context.Context, qtid QualifiedTableID, fields ...FieldVersion) error { + return nil +} + +func (s *nopVersionStore) Fields(ctx context.Context, qtid QualifiedTableID) (FieldVersions, bool, error) { + return nil, false, nil +} + +func (s *nopVersionStore) FieldVersion(ctx context.Context, qtid QualifiedTableID, field FieldName) (int, bool, error) { + return 0, true, nil +} + +func (s *nopVersionStore) FieldTables(ctx context.Context, qual TableQualifier) (TableIDs, error) { + return TableIDs{}, nil +} + +func (s *nopVersionStore) Copy(ctx context.Context) (VersionStore, error) { + return nil, nil +} diff --git a/dax/workerjob.go b/dax/workerjob.go new file mode 100644 index 000000000..490c9702e --- /dev/null +++ b/dax/workerjob.go @@ -0,0 +1,169 @@ +package dax + +import ( + "sort" + + "golang.org/x/exp/constraints" +) + +// Worker is a generic identifier used to represent a service responsible for +// doing certain jobs. In the case of dax, this is typically the Address of a +// compute or translate node. Services such as the Balancer use Workers (as +// opposed to specifically using Address) in order to remain generic, and to +// keep the business logic between services slightly less coupled. +type Worker string + +// Workers is a sortable slice of Worker. +type Workers []Worker + +func (w Workers) Len() int { return len(w) } +func (w Workers) Less(i, j int) bool { return w[i] < w[j] } +func (w Workers) Swap(i, j int) { w[i], w[j] = w[j], w[i] } + +// Job is a generic identifier used to represent a specific role assigned to a +// worker. +type Job string + +// Jobs is a slice of Job. +type Jobs []Job + +// WorkerInfo reprents a Worker and the Jobs to which it has been assigned. +type WorkerInfo struct { + ID Worker + Jobs []Job +} + +// WorkerInfos is a sortable slice of WorkerInfo. +type WorkerInfos []WorkerInfo + +func (w WorkerInfos) Len() int { return len(w) } +func (w WorkerInfos) Less(i, j int) bool { return w[i].ID < w[j].ID } +func (w WorkerInfos) Swap(i, j int) { w[i], w[j] = w[j], w[i] } + +// WorkerDiff represents the changes made to a Worker following the latest +// event. +type WorkerDiff struct { + WorkerID Worker + AddedJobs []Job + RemovedJobs []Job +} + +// Add adds w2 to w. It panics of w and w2 don't have teh same worker +// ID. Any job that is added and then removed or removed and then +// added cancels out and won't be present after add is called. +func (w *WorkerDiff) Add(w2 WorkerDiff) { + if w.WorkerID != w2.WorkerID { + panic("can't add worker diffs from different workers") + } + a1 := NewSet(w.AddedJobs...) + a2 := NewSet(w2.AddedJobs...) + r1 := NewSet(w.RemovedJobs...) + r2 := NewSet(w2.RemovedJobs...) + + // final Added is (a1 - r2) + (a2 - r1) + // this is because anything that is removed and then added, or added and then removed cancels out + added := a1.Minus(r2).Plus(a2.Minus(r1)) + + // final removed is (r1 - a2) + (r2 - a1) + removed := r1.Minus(a2).Plus(r2.Minus(a1)) + + w.AddedJobs = added.Slice() + w.RemovedJobs = removed.Slice() +} + +// WorkerDiffs is a sortable slice of WorkerDiff. +type WorkerDiffs []WorkerDiff + +func (w WorkerDiffs) Len() int { return len(w) } +func (w WorkerDiffs) Less(i, j int) bool { return w[i].WorkerID < w[j].WorkerID } +func (w WorkerDiffs) Swap(i, j int) { w[i], w[j] = w[j], w[i] } + +// Set is a set of orderable items. +type Set[K constraints.Ordered] map[K]struct{} + +func NewSet[K constraints.Ordered](stuff ...K) Set[K] { + s := make(map[K]struct{}) + for _, thing := range stuff { + s[thing] = struct{}{} + } + return Set[K](s) +} + +// Count returns the number of items in the set. +func (s Set[K]) Count() int { + return len(s) +} + +// Contains returns true if k is in the set. +func (s Set[K]) Contains(k K) bool { + _, ok := s[k] + return ok +} + +// Add adds k to the set. +func (s Set[K]) Add(k K) { + s[k] = struct{}{} +} + +// Remove removes k from the set. +func (s Set[K]) Remove(k K) { + delete(s, k) +} + +// Slice returns a slice containing each member of the set in an undefined order. +func (s Set[K]) Slice() []K { + ret := make([]K, 0, len(s)) + for k := range s { + ret = append(ret, k) + } + return ret +} + +// Copy creates a copy of the set. +func (s Set[K]) Copy() Set[K] { + ret := make(map[K]struct{}) + for k, v := range s { + ret[k] = v + } + return ret +} + +// Minus returns the a copy of s without any members which are also in s2. +func (s Set[K]) Minus(s2 Set[K]) Set[K] { + ret := make(map[K]struct{}) + for k := range s { + if _, ok := s2[k]; !ok { + ret[k] = struct{}{} + } + } + return ret +} + +// Plus returns a copy of s that also contains all members of s2. +func (s Set[K]) Plus(s2 Set[K]) Set[K] { + ret := s.Copy() + for k := range s2 { + ret[k] = struct{}{} + } + return ret +} + +// Merge adds the members of s2 to s. +func (s Set[K]) Merge(s2 Set[K]) { + for k, v := range s2 { + s[k] = v + } +} + +// Sorted returns Set[K] as a sorted slice of K. +func (s Set[K]) Sorted() []K { + js := make([]K, 0, len(s)) + for j := range s { + js = append(js, j) + } + sort.Slice(js, func(i, j int) bool { + return js[i] < js[j] + }) + + return js +} diff --git a/dax/workerjob_test.go b/dax/workerjob_test.go new file mode 100644 index 000000000..2e2726c97 --- /dev/null +++ b/dax/workerjob_test.go @@ -0,0 +1,81 @@ +package dax + +import ( + "fmt" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestWorkerDiffAdd(t *testing.T) { + tests := []struct { + w1 WorkerDiff + w2 WorkerDiff + exp WorkerDiff + }{ + { + w1: WorkerDiff{ + AddedJobs: []Job{}, + RemovedJobs: []Job{}, + }, + w2: WorkerDiff{ + AddedJobs: []Job{}, + RemovedJobs: []Job{}, + }, + exp: WorkerDiff{ + AddedJobs: []Job{}, + RemovedJobs: []Job{}, + }, + }, + { + w1: WorkerDiff{ + AddedJobs: []Job{"a", "b"}, + RemovedJobs: []Job{"c", "d"}, + }, + w2: WorkerDiff{ + AddedJobs: []Job{"c"}, + RemovedJobs: []Job{"a", "z"}, + }, + exp: WorkerDiff{ + AddedJobs: []Job{"b"}, + RemovedJobs: []Job{"d", "z"}, + }, + }, + { + w1: WorkerDiff{ + AddedJobs: []Job{"a", "b"}, + RemovedJobs: []Job{}, + }, + w2: WorkerDiff{ + AddedJobs: []Job{"c", "d"}, + RemovedJobs: []Job{}, + }, + exp: WorkerDiff{ + AddedJobs: []Job{"a", "b", "c", "d"}, + RemovedJobs: []Job{}, + }, + }, + { + w1: WorkerDiff{ + AddedJobs: []Job{}, + RemovedJobs: []Job{"a", "b"}, + }, + w2: WorkerDiff{ + AddedJobs: []Job{}, + RemovedJobs: []Job{"c", "d"}, + }, + exp: WorkerDiff{ + AddedJobs: []Job{}, + RemovedJobs: []Job{"a", "b", "c", "d"}, + }, + }, + } + + for i, tst := range tests { + t.Run(fmt.Sprintf("workerdiff: %d", i), func(t *testing.T) { + tst.w1.Add(tst.w2) + assert.ElementsMatch(t, tst.exp.AddedJobs, tst.w1.AddedJobs) + assert.ElementsMatch(t, tst.exp.RemovedJobs, tst.w1.RemovedJobs) + }) + } +} diff --git a/dax/writelogger/api/openapi.yaml b/dax/writelogger/api/openapi.yaml new file mode 100644 index 000000000..d82fe6820 --- /dev/null +++ b/dax/writelogger/api/openapi.yaml @@ -0,0 +1,113 @@ +openapi: 3.0.3 + +info: + title: WriteLogger + description: The alpha implementation of the WriteLogger interface. + version: 0.0.0 + +paths: + /writelogger/health: + get: + summary: Health check endpoint. + description: Provides an endpoint to check the overall health of the WriteLogger service. + operationId: GetHealth + responses: + 200: + description: Service is healthy. + + + /writelogger/append-message: + post: + summary: Append message to WriteLogger. + description: Appends a message to a versioned bucket/key. + operationId: PostAppendMessage + requestBody: + content: + application/json: + example: + bucket: example-bucket + key: unique-key + version: 4 + message: SGVsbG8gV29ybGQ= + schema: + type: object + properties: + bucket: + type: string + key: + type: string + version: + type: integer + format: int64 + message: + type: string + format: byte + responses: + 200: + $ref: '#/components/responses/AppendMessageResponse' + + /writelogger/log-reader: + post: + summary: Read log. + description: Reads an entire log (collection of messages) at bucket/key for the given version. + operationId: PostLogReader + requestBody: + content: + application/json: + example: + bucket: example-bucket + key: unique-key + version: 4 + schema: + type: object + properties: + bucket: + type: string + key: + type: string + version: + type: integer + format: int64 + responses: + 200: + description: Bytes making up the contents of the log. + content: + text/plain: + schema: + type: string + format: byte + + /writelogger/delete-log: + post: + summary: Delete log. + description: Deletes the log at bucket/key for the given version. + operationId: PostDeleteLog + requestBody: + content: + application/json: + example: + bucket: example-bucket + key: unique-key + version: 4 + schema: + type: object + properties: + bucket: + type: string + key: + type: string + version: + type: integer + format: int64 + responses: + 200: + description: Log was deleted. + +components: + responses: + AppendMessageResponse: + description: Placeholder response. + content: + application/json: + schema: + type: object \ No newline at end of file diff --git a/dax/writelogger/client/client.go b/dax/writelogger/client/client.go new file mode 100644 index 000000000..5e66deb73 --- /dev/null +++ b/dax/writelogger/client/client.go @@ -0,0 +1,145 @@ +// Package client contains an http implementation of the WriteLogger client. +package client + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "net/http" + + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/errors" +) + +const defaultScheme = "http" + +// WriteLogger is a client for the WriteLogger API methods. +type WriteLogger struct { + address dax.Address +} + +func New(address dax.Address) *WriteLogger { + return &WriteLogger{ + address: address, + } +} + +func (w *WriteLogger) AppendMessage(bucket string, key string, version int, msg []byte) error { + url := fmt.Sprintf("%s/writelogger/append-message", w.address.WithScheme(defaultScheme)) + + req := &AppendMessageRequest{ + Bucket: bucket, + Key: key, + Version: version, + Message: msg, + } + + // Encode the request. + postBody, err := json.Marshal(req) + if err != nil { + return errors.Wrap(err, "marshalling post request") + } + requestBody := bytes.NewBuffer(postBody) + + // Post the request. + resp, err := http.Post(url, "application/json", requestBody) + if err != nil { + return errors.Wrap(err, "posting append-message request") + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(resp.Body) + return errors.Errorf("status code: %d: %s", resp.StatusCode, b) + } + + var isr *AppendMessageResponse + if err := json.NewDecoder(resp.Body).Decode(&isr); err != nil { + return errors.Wrap(err, "reading response body") + } + + return nil +} + +type AppendMessageRequest struct { + Bucket string `json:"bucket"` + Key string `json:"key"` + Version int `json:"version"` + Message []byte `json:"message"` +} +type AppendMessageResponse struct{} + +func (w *WriteLogger) LogReader(bucket string, key string, version int) (io.Reader, io.Closer, error) { + url := fmt.Sprintf("%s/writelogger/log-reader", w.address.WithScheme(defaultScheme)) + + req := &LogReaderRequest{ + Bucket: bucket, + Version: version, + Key: key, + } + + // Encode the request. + postBody, err := json.Marshal(req) + if err != nil { + return nil, nil, errors.Wrap(err, "marshalling post request") + } + requestBody := bytes.NewBuffer(postBody) + + // Post the request. + resp, err := http.Post(url, "application/json", requestBody) + if err != nil { + return nil, nil, errors.Wrap(err, "posting log-reader request") + } + + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(resp.Body) + defer resp.Body.Close() + return nil, nil, errors.Errorf("status code: %d: %s", resp.StatusCode, b) + } + + return resp.Body, resp.Body, nil +} + +type LogReaderRequest struct { + Bucket string `json:"bucket"` + Version int `json:"version"` + Key string `json:"key"` +} + +func (w *WriteLogger) DeleteLog(bucket string, key string, version int) error { + url := fmt.Sprintf("%s/writelogger/delete-log", w.address.WithScheme(defaultScheme)) + + req := &DeleteLogRequest{ + Bucket: bucket, + Version: version, + Key: key, + } + + // Encode the request. + postBody, err := json.Marshal(req) + if err != nil { + return errors.Wrap(err, "marshalling post request") + } + requestBody := bytes.NewBuffer(postBody) + + // Post the request. + resp, err := http.Post(url, "application/json", requestBody) + if err != nil { + return errors.Wrap(err, "posting log-reader request") + } + + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(resp.Body) + defer resp.Body.Close() + return errors.Errorf("status code: %d: %s", resp.StatusCode, b) + } + + return nil +} + +type DeleteLogRequest struct { + Bucket string `json:"bucket"` + Version int `json:"version"` + Key string `json:"key"` +} diff --git a/dax/writelogger/config.go b/dax/writelogger/config.go new file mode 100644 index 000000000..6c8a5d364 --- /dev/null +++ b/dax/writelogger/config.go @@ -0,0 +1,8 @@ +package writelogger + +import "github.com/molecula/featurebase/v3/logger" + +type Config struct { + DataDir string `toml:"data-dir"` + Logger logger.Logger `toml:"-"` +} diff --git a/dax/writelogger/http/handler.go b/dax/writelogger/http/handler.go new file mode 100644 index 000000000..ea1f93e61 --- /dev/null +++ b/dax/writelogger/http/handler.go @@ -0,0 +1,121 @@ +package http + +import ( + "encoding/json" + "io" + "net/http" + + "github.com/gorilla/mux" + "github.com/molecula/featurebase/v3/dax/writelogger" + "github.com/molecula/featurebase/v3/logger" +) + +func Handler(w *writelogger.WriteLogger, logger logger.Logger) http.Handler { + svr := &server{ + writeLogger: w, + logger: logger, + } + + router := mux.NewRouter() + router.HandleFunc("/health", svr.getHealth).Methods("GET").Name("GetHealth") + router.HandleFunc("/append-message", svr.postAppendMessage).Methods("POST").Name("PostAppendMessage") + router.HandleFunc("/log-reader", svr.postLogReader).Methods("POST").Name("PostLogReader") + router.HandleFunc("/delete-log", svr.postDeleteLog).Methods("POST").Name("PostDeleteLog") + return router +} + +type server struct { + writeLogger *writelogger.WriteLogger + logger logger.Logger +} + +// GET /health +func (s *server) getHealth(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) +} + +// POST /append-message +func (s *server) postAppendMessage(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + req := AppendMessageRequest{} + if err := json.NewDecoder(body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + err := s.writeLogger.AppendMessage(req.Bucket, req.Key, req.Version, req.Message) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + resp := AppendMessageResponse{} + if err := json.NewEncoder(w).Encode(resp); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } +} + +type AppendMessageRequest struct { + Bucket string `json:"bucket"` + Key string `json:"key"` + Version int `json:"version"` + Message []byte `json:"message"` +} + +type AppendMessageResponse struct{} + +// POST /log-reader +func (s *server) postLogReader(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + req := LogReaderRequest{} + if err := json.NewDecoder(body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + reader, closer, err := s.writeLogger.LogReader(req.Bucket, req.Key, req.Version) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + defer closer.Close() + + if _, err := io.Copy(w, reader); err != nil { + s.logger.Printf("error streaming log data: %s", err) + } +} + +type LogReaderRequest struct { + Bucket string `json:"bucket"` + Version int `json:"version"` + Key string `json:"key"` +} + +// POST /delete-log +func (s *server) postDeleteLog(w http.ResponseWriter, r *http.Request) { + body := r.Body + defer body.Close() + + req := DeleteLogRequest{} + if err := json.NewDecoder(body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + if err := s.writeLogger.DeleteLog(req.Bucket, req.Key, req.Version); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } +} + +type DeleteLogRequest struct { + Bucket string `json:"bucket"` + Version int `json:"version"` + Key string `json:"key"` +} diff --git a/dax/writelogger/writelogger.go b/dax/writelogger/writelogger.go new file mode 100644 index 000000000..70b3e0530 --- /dev/null +++ b/dax/writelogger/writelogger.go @@ -0,0 +1,127 @@ +// Package writelogger provides the writelogger structs. +package writelogger + +import ( + "fmt" + "io" + "io/fs" + "os" + "path" + "sync" + + "github.com/molecula/featurebase/v3/errors" + "github.com/molecula/featurebase/v3/logger" +) + +type WriteLogger struct { + mu sync.RWMutex + + dataDir string + logFiles map[string]*os.File + + logger logger.Logger +} + +func New(cfg Config) *WriteLogger { + return &WriteLogger{ + dataDir: cfg.DataDir, + logFiles: make(map[string]*os.File), + logger: logger.NopLogger, + } +} + +// SetLogger sets the logger used for logging messages. Note, this is not the +// same "logger" that the WriteLogger represents, which logs data writes. +func (w *WriteLogger) SetLogger(l logger.Logger) { + w.logger = l +} + +func (w *WriteLogger) AppendMessage(bucket string, key string, version int, message []byte) error { + fKey := fullKey(bucket, key, version) + logFile, err := w.logFileByKey(fKey) + if err != nil { + return errors.Wrapf(err, "getting log file by key: %s", fKey) + } + + logFile.Write(append(message, "\n"...)) + logFile.Sync() + + return nil +} + +func (w *WriteLogger) LogReader(bucket string, key string, version int) (io.Reader, io.Closer, error) { + _, filePath := w.paths(fullKey(bucket, key, version)) + + f, err := os.Open(filePath) + if err != nil { + if e, ok := err.(*fs.PathError); ok { + return nil, nil, e + } + return nil, nil, err + } + w.logger.Debugf("WriteLogger LogReader file: %s", f.Name()) + + return f, f, nil +} + +func (w *WriteLogger) DeleteLog(bucket string, key string, version int) error { + w.mu.Lock() + defer w.mu.Unlock() + + fullKey := fullKey(bucket, key, version) + + f, ok := w.logFiles[fullKey] + if !ok { + return nil + } + + // Close the log file. + if err := f.Close(); err != nil { + return errors.Wrap(err, "closing log file") + } + + // Remove the log file. + return os.Remove(f.Name()) +} + +// paths takes a key and returns the full file path (including the root data +// directory) as well as the full directory path (i.e. the file path without the +// file portion). +func (w *WriteLogger) paths(key string) (string, string) { + filePath := path.Join(w.dataDir, key) + dirPath, _ := path.Split(filePath) + return dirPath, filePath +} + +// logFileByKey returns a pointer to the file specified by key. If the file does +// not exist, the file is created (along with any directories in which the file +// is nested). +func (w *WriteLogger) logFileByKey(key string) (*os.File, error) { + w.mu.Lock() + defer w.mu.Unlock() + + if f, ok := w.logFiles[key]; ok { + return f, nil + } + + dirPath, filePath := w.paths(key) + + // make directories + if err := os.MkdirAll(dirPath, 0777); err != nil { + return nil, errors.Wrapf(err, "making direcory: %s", dirPath) + } + + // open log file + f, err := os.OpenFile(filePath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644) + if err != nil { + return nil, errors.Wrapf(err, "opening file: %s", filePath) + } + w.logFiles[key] = f + + return f, nil +} + +// fullKey returns the full file key including the bucket and version. +func fullKey(bucket string, key string, version int) string { + return path.Join(bucket, key, fmt.Sprintf("%d", version)) +} diff --git a/dax/writelogger/writelogger_test.go b/dax/writelogger/writelogger_test.go new file mode 100644 index 000000000..7cfb6dfd4 --- /dev/null +++ b/dax/writelogger/writelogger_test.go @@ -0,0 +1,73 @@ +package writelogger_test + +import ( + "encoding/json" + "fmt" + "io" + "os" + "path" + "testing" + + "github.com/molecula/featurebase/v3/dax/writelogger" + "github.com/stretchr/testify/assert" +) + +func TestWriteLogger(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "testWriteLogger-*") + assert.NoError(t, err) + + // Remove the temp directory. + defer func() { + os.RemoveAll(tmpDir) + }() + + t.Run("Basic", func(t *testing.T) { + type payload struct { + Foo string `json:"foo"` + Bar int `json:"bar"` + } + + cfg := writelogger.Config{ + DataDir: tmpDir, + } + wl := writelogger.New(cfg) + + table := "tbl" + partition := 1 + version := 0 + key := "keys" + + msg1 := payload{ + Foo: "message 1", + Bar: 88, + } + + // Write the message. + msg, err := json.Marshal(msg1) + assert.NoError(t, err) + + err = wl.AppendMessage(bucket(table, partition), key, version, msg) + assert.NoError(t, err) + + // Read the message. + reader, closer, err := wl.LogReader(bucket(table, partition), key, version) + assert.NoError(t, err) + defer closer.Close() + + buf, err := io.ReadAll(reader) + assert.NoError(t, err) + + var out payload + + err = json.Unmarshal(buf, &out) + assert.NoError(t, err) + + assert.Equal(t, msg1.Foo, out.Foo) + assert.Equal(t, msg1.Bar, out.Bar) + }) +} + +func bucket(table string, partition int) string { + return path.Join(table, fmt.Sprintf("%d", partition)) + +} diff --git a/errors/errors.go b/errors/errors.go index 534760bff..230788aaf 100644 --- a/errors/errors.go +++ b/errors/errors.go @@ -1,4 +1,4 @@ -// Package errors wraps pkg/errors and includes some custom featurs such as +// Package errors wraps pkg/errors and includes some custom features such as // error codes. package errors @@ -83,3 +83,7 @@ func (ce codedError) Is(err error) bool { } return false } + +const ( + ErrUncoded Code = "Uncoded" +) diff --git a/errors/errors_test.go b/errors/errors_test.go index 6e0b59fc6..e8c5e8782 100644 --- a/errors/errors_test.go +++ b/errors/errors_test.go @@ -9,27 +9,8 @@ import ( ) func TestErrors(t *testing.T) { - - var errUncoded errors.Code = "TestErrUncoded" - var errFieldNotFound errors.Code = "TestErrFieldNotFound" - var errTableNotFound errors.Code = "TestErrTableNotFound" - - newErrFieldNotFound := func(fld string) error { - return errors.New( - errFieldNotFound, - fmt.Sprintf("field not found '%s'", fld), - ) - } - - newErrTableNotFound := func(tbl string) error { - return errors.New( - errTableNotFound, - fmt.Sprintf("table not found '%s'", tbl), - ) - } - t.Run("Is", func(t *testing.T) { - uncoded := errors.New(errUncoded, "uncoded error") + uncoded := newUncoded("uncoded error") fnf := newErrFieldNotFound("fld") tnf := newErrTableNotFound("tbl") fnfCustom := errors.New(errFieldNotFound, "custom field message") @@ -79,3 +60,32 @@ func TestErrors(t *testing.T) { } }) } + +// Test error codes. + +const ( + errUncoded errors.Code = "Uncoded" + errFieldNotFound errors.Code = "FieldNotFound" + errTableNotFound errors.Code = "TableNotFound" +) + +func newUncoded(message string) error { + return errors.New( + errUncoded, + message, + ) +} + +func newErrFieldNotFound(field string) error { + return errors.New( + errFieldNotFound, + "field not found: "+field, + ) +} + +func newErrTableNotFound(table string) error { + return errors.New( + errTableNotFound, + "table not found: "+table, + ) +} diff --git a/executor.go b/executor.go index 9ccb0ee0e..9eee46c8f 100644 --- a/executor.go +++ b/executor.go @@ -350,9 +350,9 @@ func (e *executor) handlePreCalls(ctx context.Context, qcx *Qcx, index string, c idx := c.Args["valueidx"].(int64) if idx >= 0 && idx < int64(len(opt.EmbeddedData)) { row := opt.EmbeddedData[idx] - c.Precomputed = make(map[uint64]interface{}, len(row.segments)) - for _, segment := range row.segments { - c.Precomputed[segment.shard] = &Row{segments: []rowSegment{segment}} + c.Precomputed = make(map[uint64]interface{}, len(row.Segments)) + for _, segment := range row.Segments { + c.Precomputed[segment.shard] = &Row{Segments: []RowSegment{segment}} } } else { return fmt.Errorf("no precomputed data! index %d, len %d", idx, len(opt.EmbeddedData)) @@ -425,9 +425,9 @@ func (e *executor) handlePreCalls(ctx context.Context, qcx *Qcx, index string, c opt.EmbeddedData = append(opt.EmbeddedData, row) // and stash a copy locally, so local calls can use it if row != nil { - c.Precomputed = make(map[uint64]interface{}, len(row.segments)) - for _, segment := range row.segments { - c.Precomputed[segment.shard] = &Row{segments: []rowSegment{segment}} + c.Precomputed = make(map[uint64]interface{}, len(row.Segments)) + for _, segment := range row.Segments { + c.Precomputed[segment.shard] = &Row{Segments: []RowSegment{segment}} } } return nil @@ -582,7 +582,7 @@ func (e *executor) execute(ctx context.Context, qcx *Qcx, index string, q *pql.Q } if vc, ok := v.(ValCount); ok { - vc.cleanup() + vc.Cleanup() v = vc } @@ -606,7 +606,7 @@ func (e *executor) execute(ctx context.Context, qcx *Qcx, index string, q *pql.Q // functions may return both integer and the interpreted value, but we // don't want to pass that all the way back to the client, so we // remove it here. -func (vc *ValCount) cleanup() { +func (vc *ValCount) Cleanup() { if vc.Val != 0 && (vc.FloatVal != 0 || !vc.TimestampVal.IsZero() || vc.DecimalVal != nil) { vc.Val = 0 } @@ -1029,8 +1029,8 @@ func (e *executor) executeLimitCall(ctx context.Context, qcx *Qcx, index string, if offset != 0 { i := 0 var leadingBits []uint64 - for i < len(result.segments) && offset > 0 { - seg := result.segments[i] + for i < len(result.Segments) && offset > 0 { + seg := result.Segments[i] count := seg.Count() if count > offset { data := seg.Columns() @@ -1044,14 +1044,14 @@ func (e *executor) executeLimitCall(ctx context.Context, qcx *Qcx, index string, i++ } row := NewRow(leadingBits...) - row.Merge(&Row{segments: result.segments[i:]}) + row.Merge(&Row{Segments: result.Segments[i:]}) result = row } if limit < result.Count() { i := 0 var trailingBits []uint64 - for i < len(result.segments) && limit > 0 { - seg := result.segments[i] + for i < len(result.Segments) && limit > 0 { + seg := result.Segments[i] count := seg.Count() if count > limit { data := seg.Columns() @@ -1064,7 +1064,7 @@ func (e *executor) executeLimitCall(ctx context.Context, qcx *Qcx, index string, i++ } row := NewRow(trailingBits...) - row.Merge(&Row{segments: result.segments[:i]}) + row.Merge(&Row{Segments: result.Segments[:i]}) result = row } @@ -1109,7 +1109,7 @@ func (e *executor) executeSum(ctx context.Context, qcx *Qcx, index string, c *pq // Merge returned results at coordinating node. reduceFn := func(ctx context.Context, prev, v interface{}) interface{} { other, _ := prev.(ValCount) - return other.add(v.(ValCount)) + return other.Add(v.(ValCount)) } result, err := e.mapReduce(ctx, index, shards, c, opt, mapFn, reduceFn) @@ -1165,7 +1165,7 @@ func (e *executor) executeDistinct(ctx context.Context, qcx *Qcx, index string, } switch other := prev.(type) { case SignedRow: - return other.union(v.(SignedRow)) + return other.Union(v.(SignedRow)) case *Row: if other == nil { return v @@ -1188,7 +1188,7 @@ func (e *executor) executeDistinct(ctx context.Context, qcx *Qcx, index string, } if other, ok := result.(SignedRow); ok { - other.field = field + other.Field = field } return result, nil } @@ -1214,7 +1214,7 @@ func (e *executor) executeMin(ctx context.Context, qcx *Qcx, index string, c *pq // Merge returned results at coordinating node. reduceFn := func(ctx context.Context, prev, v interface{}) interface{} { other, _ := prev.(ValCount) - return other.smaller(v.(ValCount)) + return other.Smaller(v.(ValCount)) } result, err := e.mapReduce(ctx, index, shards, c, opt, mapFn, reduceFn) @@ -1250,7 +1250,7 @@ func (e *executor) executeMax(ctx context.Context, qcx *Qcx, index string, c *pq // Merge returned results at coordinating node. reduceFn := func(ctx context.Context, prev, v interface{}) interface{} { other, _ := prev.(ValCount) - return other.larger(v.(ValCount)) + return other.Larger(v.(ValCount)) } result, err := e.mapReduce(ctx, index, shards, c, opt, mapFn, reduceFn) @@ -1598,8 +1598,8 @@ func (e *executor) executeDistinctShard(ctx context.Context, qcx *Qcx, index str return result, errors.Wrap(err, "executing bitmap call") } filter = row - if filter != nil && len(filter.segments) > 0 { - filterBitmap = filter.segments[0].data + if filter != nil && len(filter.Segments) > 0 { + filterBitmap = filter.Segments[0].data } // if we had a filter to consider, but it came back empty, we // can go ahead and save time by returning the empty results, @@ -2113,7 +2113,7 @@ func (e *executor) executeTopK(ctx context.Context, qcx *Qcx, index string, c *p reduceFn := func(ctx context.Context, prev, v interface{}) interface{} { x, _ := prev.([]*Row) y, _ := v.([]*Row) - return ([]*Row)(addBSI(x, y)) + return ([]*Row)(AddBSI(x, y)) } other, err := e.mapReduce(ctx, index, shards, c, opt, mapFn, reduceFn) @@ -2137,7 +2137,7 @@ func (e *executor) executeTopK(ctx context.Context, qcx *Qcx, index string, c *p } var dst []Pair - bsiData(results).pivotDescending(NewRow().Union(results...), 0, limit, nil, func(count uint64, ids ...uint64) { + BSIData(results).PivotDescending(NewRow().Union(results...), 0, limit, nil, func(count uint64, ids ...uint64) { for _, id := range ids { dst = append(dst, Pair{ ID: id, @@ -2280,7 +2280,7 @@ func (e *executor) executeTopKShardTime(ctx context.Context, tx Tx, filter *Row, // topKFragments builds a perpendicular BSI bitmap from fragments. // The fragments are expected to be from set fields. -func topKFragments(ctx context.Context, tx Tx, filter *Row, fragments ...*fragment) (bsiData, error) { +func topKFragments(ctx context.Context, tx Tx, filter *Row, fragments ...*fragment) (BSIData, error) { // Acquire fragment container iterators. iters := make([]roaring.ContainerIterator, len(fragments)) for i, f := range fragments { @@ -2450,7 +2450,7 @@ func (h mergeratorHeap) minHeapify() { // doTopK uses a raw Pilosa matrix to produce a perpendicular BSI bitmap. // It will apply a row filter if one is provided. -func doTopK(ctx context.Context, it roaring.ContainerIterator, filter *topKFilter) (bsiData, error) { +func doTopK(ctx context.Context, it roaring.ContainerIterator, filter *topKFilter) (BSIData, error) { row := ^uint64(0) var count uint64 @@ -2499,7 +2499,7 @@ type topKFilter [ShardWidth >> 16]*roaring.Container // fill the filter with the contents of a Row. func (f *topKFilter) fill(row *Row) { - for _, s := range row.segments { + for _, s := range row.Segments { it, _ := s.data.Containers.Iterator(0) f.fillIt(it) } @@ -2727,12 +2727,12 @@ func (e *executor) executeDifferenceShard(ctx context.Context, qcx *Qcx, index s type RowIdentifiers struct { Rows []uint64 `json:"rows"` Keys []string `json:"keys,omitempty"` - field string + Field string } func (r *RowIdentifiers) Clone() (clone *RowIdentifiers) { clone = &RowIdentifiers{ - field: r.field, + Field: r.Field, } if r.Rows != nil { clone.Rows = make([]uint64, len(r.Rows)) @@ -2759,7 +2759,7 @@ func (r RowIdentifiers) ToTable() (*proto.TableResponse, error) { // ToRows implements the ToRowser interface. func (r RowIdentifiers) ToRows(callback func(*proto.RowResponse) error) error { if len(r.Keys) > 0 { - ci := []*proto.ColumnInfo{{Name: r.Field(), Datatype: "string"}} + ci := []*proto.ColumnInfo{{Name: r.Field, Datatype: "string"}} for _, key := range r.Keys { if err := callback(&proto.RowResponse{ Headers: ci, @@ -2772,7 +2772,7 @@ func (r RowIdentifiers) ToRows(callback func(*proto.RowResponse) error) error { ci = nil } } else { - ci := []*proto.ColumnInfo{{Name: r.Field(), Datatype: "uint64"}} + ci := []*proto.ColumnInfo{{Name: r.Field, Datatype: "uint64"}} for _, id := range r.Rows { if err := callback(&proto.RowResponse{ Headers: ci, @@ -2788,18 +2788,13 @@ func (r RowIdentifiers) ToRows(callback func(*proto.RowResponse) error) error { return nil } -// Field returns the field name associated to the row. -func (r *RowIdentifiers) Field() string { - return r.field -} - // RowIDs is a query return type for just uint64 row ids. // It should only be used internally (since RowIdentifiers // is the external return type), but it is exported because // the proto package needs access to it. type RowIDs []uint64 -func (r RowIDs) merge(other RowIDs, limit int) RowIDs { +func (r RowIDs) Merge(other RowIDs, limit int) RowIDs { i, j := 0, 0 result := make(RowIDs, 0) for i < len(r) && j < len(other) && len(result) < limit { @@ -3151,7 +3146,7 @@ func (e *executor) executeGroupBy(ctx context.Context, qcx *Qcx, index string, c for subj, cond := range having.Args { switch subj { case "count", "sum": - results = applyConditionToGroupCounts(results, subj, cond.(*pql.Condition)) + results = ApplyConditionToGroupCounts(results, subj, cond.(*pql.Condition)) default: return nil, errors.New("Condition() only supports count or sum") } @@ -3654,10 +3649,10 @@ func (g GroupCount) satisfiesCondition(subj string, cond *pql.Condition) bool { return false } -// applyConditionToGroupCounts filters the contents of gcs according +// ApplyConditionToGroupCounts filters the contents of gcs according // to the condition. Currently, `count` and `sum` are the only // fields supported. -func applyConditionToGroupCounts(gcs []GroupCount, subj string, cond *pql.Condition) []GroupCount { +func ApplyConditionToGroupCounts(gcs []GroupCount, subj string, cond *pql.Condition) []GroupCount { var i int for _, gc := range gcs { if !gc.satisfiesCondition(subj, cond) { @@ -3788,7 +3783,7 @@ func (e *executor) executeRows(ctx context.Context, qcx *Qcx, index string, c *p if err := ctx.Err(); err != nil { return err } - return other.merge(v.(RowIDs), limit) + return other.Merge(v.(RowIDs), limit) } // Get full result set. other, err := e.mapReduce(ctx, index, shards, c, opt, mapFn, reduceFn) @@ -3926,7 +3921,7 @@ func (e *executor) executeRowsShard(ctx context.Context, qcx *Qcx, index string, if err != nil { return nil, err } - rowIDs = rowIDs.merge(viewRows, limit) + rowIDs = rowIDs.Merge(viewRows, limit) } return rowIDs, nil @@ -6144,14 +6139,14 @@ func makeEmbeddedDataForShards(allRows []*Row, shards []uint64) []*Row { } newRows := make([]*Row, len(allRows)) for i, row := range allRows { - if row == nil || len(row.segments) == 0 { + if row == nil || len(row.Segments) == 0 { continue } if row.NoSplit { newRows[i] = row continue } - segments := row.segments + segments := row.Segments segmentIndex := 0 newRows[i] = &Row{ Index: row.Index, @@ -6166,7 +6161,7 @@ func makeEmbeddedDataForShards(allRows []*Row, shards []uint64) []*Row { break } if segments[segmentIndex].shard == shard { - newRows[i].segments = append(newRows[i].segments, segments[segmentIndex]) + newRows[i].Segments = append(newRows[i].Segments, segments[segmentIndex]) segmentIndex++ if segmentIndex >= len(segments) { // no more segments, we're done @@ -7231,7 +7226,7 @@ func (e *executor) collectResultIDs(index string, idx *Index, call *pql.Call, re return errors.Wrap(err, "determining how to translate") } if strategy == byCurrentIndex { - for _, segment := range result.Segments() { + for _, segment := range result.Segments { for _, col := range segment.Columns() { idSet[col] = struct{}{} } @@ -7268,7 +7263,7 @@ func (e *executor) translateResult(ctx context.Context, index string, idx *Index switch strategy { case byCurrentIndex: other := &Row{} - for _, segment := range result.Segments() { + for _, segment := range result.Segments { for _, col := range segment.Columns() { other.Keys = append(other.Keys, idSet[col]) } @@ -7285,7 +7280,7 @@ func (e *executor) translateResult(ctx context.Context, index string, idx *Index if idx == nil { return nil, errors.Errorf("foreign index %s not found for field %s in index %s", rowField.ForeignIndex(), rowField.Name(), rowField.Index()) } - for _, segment := range result.Segments() { + for _, segment := range result.Segments { keys, err := e.Cluster.translateIndexIDs(context.Background(), rowField.ForeignIndex(), segment.Columns()) if err != nil { return nil, errors.Wrap(err, "translating index ids") @@ -7294,7 +7289,7 @@ func (e *executor) translateResult(ctx context.Context, index string, idx *Index } case byRowIndex: - for _, segment := range result.Segments() { + for _, segment := range result.Segments { keys, err := e.Cluster.translateIndexIDs(context.Background(), rowIdx.Name(), segment.Columns()) if err != nil { return nil, errors.Wrap(err, "translating index ids") @@ -7326,7 +7321,7 @@ func (e *executor) translateResult(ctx context.Context, index string, idx *Index return &SignedRow{Pos: &Row{}}, nil } other := &Row{} - for _, segment := range rslt.Segments() { + for _, segment := range rslt.Segments { keys, err := e.Cluster.translateIndexIDs(context.Background(), field.ForeignIndex(), segment.Columns()) if err != nil { return nil, errors.Wrap(err, "translating index ids") @@ -7478,7 +7473,7 @@ func (e *executor) translateResult(ctx context.Context, index string, idx *Index } other := RowIdentifiers{ - field: fieldName, + Field: fieldName, } if field := idx.Field(fieldName); field == nil { @@ -7825,25 +7820,20 @@ func needsShards(call *pql.Call) bool { // SignedRow represents a signed *Row with two (neg/pos) *Rows. type SignedRow struct { - Neg *Row `json:"neg"` - Pos *Row `json:"pos"` - field string + Neg *Row `json:"neg"` + Pos *Row `json:"pos"` + Field string `json:"-"` } func (s *SignedRow) Clone() (r *SignedRow) { r = &SignedRow{ Neg: s.Neg.Clone(), // Row.Clone() returns nil for nil. Pos: s.Pos.Clone(), - field: s.field, + Field: s.Field, } return } -// Field returns the field name associated to the signed row. -func (s *SignedRow) Field() string { - return s.field -} - // ToTable implements the ToTabler interface. func (s SignedRow) ToTable() (*proto.TableResponse, error) { var n uint64 @@ -7858,7 +7848,7 @@ func (s SignedRow) ToTable() (*proto.TableResponse, error) { // ToRows implements the ToRowser interface. func (s SignedRow) ToRows(callback func(*proto.RowResponse) error) error { - ci := []*proto.ColumnInfo{{Name: s.Field(), Datatype: "int64"}} + ci := []*proto.ColumnInfo{{Name: s.Field, Datatype: "int64"}} if s.Neg != nil { negs := s.Neg.Columns() for i := len(negs) - 1; i >= 0; i-- { @@ -7924,7 +7914,7 @@ func toInt64(n uint64) (int64, error) { return int64(n), nil } -func (sr *SignedRow) union(other SignedRow) SignedRow { +func (sr *SignedRow) Union(other SignedRow) SignedRow { ret := SignedRow{&Row{}, &Row{}, ""} // merge in sr @@ -8042,7 +8032,7 @@ func (v ValCount) ToRows(callback func(*proto.RowResponse) error) error { return nil } -func (vc *ValCount) add(other ValCount) ValCount { +func (vc *ValCount) Add(other ValCount) ValCount { return ValCount{ Val: vc.Val + other.Val, Count: vc.Count + other.Count, @@ -8050,7 +8040,7 @@ func (vc *ValCount) add(other ValCount) ValCount { } // smaller returns the smaller of the two ValCounts. -func (vc *ValCount) smaller(other ValCount) ValCount { +func (vc *ValCount) Smaller(other ValCount) ValCount { if vc.DecimalVal != nil || other.DecimalVal != nil { return vc.decimalSmaller(other) } else if vc.FloatVal != 0 || other.FloatVal != 0 { @@ -8130,7 +8120,7 @@ func (vc *ValCount) floatSmaller(other ValCount) ValCount { } // larger returns the larger of the two ValCounts. -func (vc *ValCount) larger(other ValCount) ValCount { +func (vc *ValCount) Larger(other ValCount) ValCount { if vc.DecimalVal != nil || other.DecimalVal != nil { return vc.decimalLarger(other) } else if vc.FloatVal != 0 || other.FloatVal != 0 { @@ -8724,10 +8714,10 @@ func (e *executor) executeDeleteRecordFromShard(ctx context.Context, index strin if err != nil { return false, err } - if len(row.segments) == 0 { + if len(row.Segments) == 0 { return } - columns := row.segments[0].data + columns := row.Segments[0].data if columns.Count() == 0 { return } @@ -8877,14 +8867,14 @@ func DeleteRowsWithOutKeysFlow(ctx context.Context, columns *roaring.Bitmap, idx } func DeleteRowsWithFlow(ctx context.Context, src *Row, idx *Index, shard uint64, normalFlow bool) (change bool, err error) { - if len(src.segments) == 0 { // nothing to remove + if len(src.Segments) == 0 { //nothing to remove return false, nil } - columns := src.segments[0].data // should only be one segment + columns := src.Segments[0].data //should only be one segment if columns.Count() == 0 { return false, nil } - bits := src.segments[0].data.Slice() + bits := src.Segments[0].data.Slice() min := func(a, b int) int { if a <= b { return a diff --git a/executor_internal_test.go b/executor_internal_test.go index 2d20d553e..8f483e4d0 100644 --- a/executor_internal_test.go +++ b/executor_internal_test.go @@ -322,12 +322,12 @@ func TestValCountComparisons(t *testing.T) { for i, test := range tests { t.Run(test.name+strconv.Itoa(i), func(t *testing.T) { - gotLarger := test.vc.larger(test.other) + gotLarger := test.vc.Larger(test.other) if gotLarger != test.expLarger { t.Fatalf("larger failed, expected:\n%+v\ngot:\n%+v", test.expLarger, gotLarger) } - gotSmaller := test.vc.smaller(test.other) + gotSmaller := test.vc.Smaller(test.other) if gotSmaller != test.expSmaller { t.Fatalf("smaller failed, expected:\n%+v\ngot:\n%+v", test.expSmaller, gotSmaller) } diff --git a/fbcloud/auth.go b/fbcloud/auth.go new file mode 100644 index 000000000..c1c59f6ef --- /dev/null +++ b/fbcloud/auth.go @@ -0,0 +1,79 @@ +package fbcloud + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "net/http" + + "github.com/pkg/errors" +) + +const ( + authFlow = "USER_PASSWORD_AUTH" + cognitoURLTemplate = "https://cognito-idp.%s.amazonaws.com" +) + +type cognitoParameters struct { + Email string `json:"USERNAME"` + Password string `json:"PASSWORD"` +} + +type cognitoAuthRequest struct { + AuthParameters cognitoParameters `json:"AuthParameters"` + AuthFlow string `json:"AuthFlow"` + AppClientId string `json:"ClientId"` +} + +type cognitoAuthResult struct { + IdToken string `json:"IdToken"` +} + +type cognitoAuthResponse struct { + Result cognitoAuthResult `json:"AuthenticationResult"` +} + +func authenticate(clientID, region, email, password string) (string, error) { + authPayload := cognitoAuthRequest{ + AuthParameters: cognitoParameters{ + Email: email, + Password: password, + }, + AuthFlow: authFlow, + AppClientId: clientID, + } + + data, err := json.Marshal(authPayload) + if err != nil { + return "", errors.Wrap(err, "marshaling json") + } + + url := fmt.Sprintf(cognitoURLTemplate, region) + req, err := http.NewRequest(http.MethodPost, url, bytes.NewBuffer(data)) + if err != nil { + return "", errors.Wrap(err, "creating authentication request object") + } + req.Header.Add("Content-Type", "application/x-amz-json-1.1") + req.Header.Add("X-Amz-Target", "AWSCognitoIdentityProviderService.InitiateAuth") + + resp, err := http.DefaultClient.Do(req) + if err != nil { + return "", errors.Wrap(err, "making request") + } + defer resp.Body.Close() + + fullbod, err := io.ReadAll(resp.Body) + + if resp.StatusCode != http.StatusOK || err != nil { + return "", errors.Errorf("HTTP status code=%d from Cognito authentication response. reading body: %v, body: '%s'", resp.StatusCode, err, fullbod) + } + + var auth cognitoAuthResponse + err = json.Unmarshal(fullbod, &auth) + if err != nil { + return "", errors.Wrap(err, "decoding cognito auth response") + } + + return auth.Result.IdToken, nil +} diff --git a/fbcloud/client.go b/fbcloud/client.go new file mode 100644 index 000000000..fbb9e0233 --- /dev/null +++ b/fbcloud/client.go @@ -0,0 +1,147 @@ +package fbcloud + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "time" + + featurebase "github.com/molecula/featurebase/v3" + "github.com/pkg/errors" +) + +// TokenRefreshTimeout is currently hardcoded to be just under the +// Cognito token timeout for cloud which is 15 minutes (I think I +// heard that somewhere anyway). It seems to work. +const TokenRefreshTimeout = time.Minute * 13 + +type Queryer struct { + Host string + + ClientID string + Region string + Email string + Password string + + token string + lastRefresh time.Time +} + +func (cq *Queryer) tokenRefresh() error { + token, err := authenticate(cq.ClientID, cq.Region, cq.Email, cq.Password) + if err != nil { + return errors.Wrap(err, "getting token") + } + cq.token = token + cq.lastRefresh = time.Now() + fmt.Println("refreshed auth token") + return nil +} + +type tokenizedSQL struct { + Language string `json:"language"` + Statement string `json:"statement"` +} + +// Query issues a SQL query formatted for the FeatureBase cloud query endpoint. +func (cq *Queryer) Query(org, db, sql string) (*featurebase.SQLResponse, error) { + if time.Since(cq.lastRefresh) > TokenRefreshTimeout { + if err := cq.tokenRefresh(); err != nil { + return nil, errors.Wrap(err, "refreshing token") + } + } + url := fmt.Sprintf("%s/v2/databases/%s/query", cq.Host, db) + + sqlReq := &tokenizedSQL{ + Language: "sql", + Statement: sql, + } + var buf bytes.Buffer + if err := json.NewEncoder(&buf).Encode(sqlReq); err != nil { + return nil, errors.Wrapf(err, "encoding sql request: %s", sql) + } + + client := &http.Client{ + Timeout: time.Second * 30, + } + req, err := http.NewRequest(http.MethodPost, url, &buf) + if err != nil { + return nil, errors.Wrap(err, "creating new post request") + } + req.Header.Add("Authorization", cq.token) + + var resp *http.Response + if resp, err = client.Do(req); err != nil { + return nil, errors.Wrap(err, "executing post request") + } + + fullbod, err := io.ReadAll(resp.Body) + if err != nil { + return nil, errors.Wrap(err, "reading cloud response") + } + if resp.StatusCode/100 != 2 { + return nil, errors.Errorf("unexpected status: %s, full body: '%s'", resp.Status, fullbod) + } + var cloudResp cloudResponse + if err := json.Unmarshal(fullbod, &cloudResp); err != nil { + return nil, errors.Wrapf(err, "decoding cloud response, body:\n%s", fullbod) + } + sqlResponse := cloudResp.Results + return &sqlResponse, nil +} + +// HTTPRequest can make an arbitrary http request to the host and +// tries to json unmarshal the response body into v if v is +// non-nil. This is handy for hitting cloud endpoints other than the +// query endpoint which is handled by Query. I don't think this is +// currently used, but I'd like to keep it around for debugging. +func (cq *Queryer) HTTPRequest(method, path, body string, v interface{}) ([]byte, error) { + if time.Since(cq.lastRefresh) > TokenRefreshTimeout { + if err := cq.tokenRefresh(); err != nil { + return nil, errors.Wrap(err, "refreshing token") + } + } + var bod io.Reader + if body == "" { + bod = nil + } else { + bod = strings.NewReader(body) + } + req, err := http.NewRequest(method, fmt.Sprintf("%s%s", cq.Host, path), bod) + if err != nil { + return nil, errors.Errorf("creating request: %v", err) + } + // fmt.Printf("%+v\n", req) + + req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", cq.token)) + if bod != nil { + req.Header.Add("Content-Type", "application/json") + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + return nil, errors.Errorf("doing request: %v", err) + } + bodbytes, err := io.ReadAll(resp.Body) + if err != nil { + return nil, errors.Errorf("reading response body: %v", err) + } + if resp.StatusCode/100 != 2 { + return nil, errors.Errorf("bad status: %s. body: '%s'", resp.Status, bodbytes) + } + + if v != nil { + err = json.Unmarshal(bodbytes, v) + if err != nil { + return nil, errors.Errorf("unmarshaling: %v", err) + } + } + + return bodbytes, nil +} + +type cloudResponse struct { + Results featurebase.SQLResponse `json:"results"` +} diff --git a/fragment.go b/fragment.go index 8eb2bcf6e..c00b9ec8a 100644 --- a/fragment.go +++ b/fragment.go @@ -328,7 +328,7 @@ func (f *fragment) rowFromStorage(tx Tx, rowID uint64) (*Row, error) { } row := &Row{ - segments: []rowSegment{{ + Segments: []RowSegment{{ data: data, shard: f.shard, writable: true, @@ -735,7 +735,7 @@ func (f *fragment) sum(tx Tx, filter *Row, bitDepth uint64) (sum int64, count ui // though, we want to run with no-filter, as opposed to an empty filter. var filterData *roaring.Bitmap if filter != nil { - for _, seg := range filter.segments { + for _, seg := range filter.Segments { if seg.shard == f.shard { filterData = seg.data break @@ -2524,7 +2524,7 @@ func (f *fragment) unprotectedUnionRows(ctx context.Context, tx Tx, rows []uint6 return nil, err } else { row := &Row{ - segments: []rowSegment{{ + Segments: []RowSegment{{ data: filter.Results(f.shard), shard: f.shard, writable: true, diff --git a/go.mod b/go.mod index c084ab47f..774ab02de 100644 --- a/go.mod +++ b/go.mod @@ -8,6 +8,7 @@ require ( github.com/CAFxX/gcnotifier v0.0.0-20220409005548-0153238b886a github.com/DataDog/datadog-go v4.8.3+incompatible github.com/HdrHistogram/hdrhistogram-go v1.1.2 // indirect + github.com/Microsoft/go-winio v0.5.2 // indirect github.com/alexbrainman/odbc v0.0.0-20211220213544-9c9a2e61c5e2 github.com/aws/aws-sdk-go v1.42.39 github.com/beevik/ntp v0.3.0 @@ -17,6 +18,10 @@ require ( github.com/confluentinc/confluent-kafka-go v1.9.1 github.com/davecgh/go-spew v1.1.1 github.com/denisenkom/go-mssqldb v0.11.0 + github.com/docker/distribution v2.8.1+incompatible // indirect + github.com/docker/docker v20.10.17+incompatible + github.com/docker/go-connections v0.4.0 + github.com/docker/go-units v0.4.0 // indirect github.com/felixge/fgprof v0.9.2 github.com/getsentry/sentry-go v0.13.0 github.com/glycerine/vprint v0.0.0-20200730000117-76cea49a68ea @@ -28,16 +33,16 @@ require ( github.com/golang-jwt/jwt v3.2.2+incompatible github.com/golang/protobuf v1.5.2 github.com/google/go-cmp v0.5.8 + github.com/gopherjs/gopherjs v0.0.0-20200217142428-fce0ec30dd00 // indirect github.com/gorilla/handlers v1.3.0 github.com/gorilla/mux v1.8.0 github.com/gorilla/securecookie v1.1.1 github.com/hashicorp/go-retryablehttp v0.7.1 github.com/improbable-eng/grpc-web v0.15.0 github.com/jedib0t/go-pretty v4.3.0+incompatible - github.com/jonboulle/clockwork v0.3.0 // indirect - github.com/klauspost/compress v1.15.1 // indirect github.com/lib/pq v1.10.5 github.com/molecula/apophenia v0.0.0-20190827192002-68b7a14a478b + github.com/opencontainers/image-spec v1.0.2 github.com/opentracing/opentracing-go v1.2.0 github.com/pelletier/go-toml v1.9.5 github.com/pkg/errors v0.9.1 @@ -47,6 +52,7 @@ require ( github.com/rakyll/statik v0.1.7 github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0 // indirect github.com/ricochet2200/go-disk-usage/du v0.0.0-20210707232629-ac9918953285 + github.com/rs/cors v1.8.2 // indirect github.com/satori/go.uuid v1.2.1-0.20180404165556-75cca531ea76 github.com/segmentio/kafka-go v0.4.29 github.com/shirou/gopsutil/v3 v3.22.5 @@ -63,8 +69,9 @@ require ( go.etcd.io/etcd/client/pkg/v3 v3.5.4 go.etcd.io/etcd/client/v3 v3.5.4 go.etcd.io/etcd/server/v3 v3.5.4 - golang.org/x/exp v0.0.0-20201008143054-e3b2a7f2fdc7 - golang.org/x/mod v0.5.1 + golang.org/x/exp v0.0.0-20221031165847-c99f073a8326 + golang.org/x/mod v0.6.0 + golang.org/x/net v0.1.0 // indirect golang.org/x/oauth2 v0.0.0-20220608161450-d0670ef3b1eb golang.org/x/sync v0.0.0-20220601150217-0de741cfad7f golang.org/x/time v0.0.0-20211116232009-f0f3c7e86c11 @@ -80,7 +87,7 @@ require ( github.com/PaesslerAG/gval v1.0.0 github.com/PaesslerAG/jsonpath v0.1.1 github.com/google/uuid v1.3.0 - github.com/jaffee/commandeer v0.5.0 + github.com/jaffee/commandeer v0.6.0 github.com/linkedin/goavro/v2 v2.11.1 google.golang.org/grpc v1.46.0 google.golang.org/protobuf v1.28.0 @@ -89,7 +96,6 @@ require ( require ( github.com/DataDog/datadog-go/v5 v5.1.0 // indirect github.com/DataDog/gostackparse v0.5.0 // indirect - github.com/Microsoft/go-winio v0.5.2 // indirect github.com/asaskevich/govalidator v0.0.0-20200907205600-7a23bdc65eef // indirect github.com/beorn7/perks v1.0.1 // indirect github.com/cenkalti/backoff/v4 v4.1.3 // indirect @@ -103,12 +109,14 @@ require ( github.com/go-ole/go-ole v1.2.6 // indirect github.com/go-openapi/errors v0.19.8 // indirect github.com/go-stack/stack v1.8.0 // indirect + github.com/gobwas/httphead v0.0.0-20200921212729-da3d93bc3c58 // indirect + github.com/gobwas/pool v0.2.1 // indirect + github.com/gobwas/ws v1.0.4 // indirect github.com/golang-sql/civil v0.0.0-20190719163853-cb61b32ac6fe // indirect github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b // indirect github.com/golang/snappy v0.0.4 // indirect github.com/google/btree v1.0.1 // indirect github.com/google/pprof v0.0.0-20211214055906-6f57359322fd // indirect - github.com/gopherjs/gopherjs v0.0.0-20200217142428-fce0ec30dd00 // indirect github.com/gorilla/websocket v1.5.0 // indirect github.com/grpc-ecosystem/go-grpc-middleware v1.3.0 // indirect github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0 // indirect @@ -117,22 +125,26 @@ require ( github.com/hashicorp/hcl v1.0.0 // indirect github.com/inconshreveable/mousetrap v1.0.0 // indirect github.com/jmespath/go-jmespath v0.4.0 // indirect + github.com/jonboulle/clockwork v0.2.2 // indirect github.com/json-iterator/go v1.1.12 // indirect + github.com/klauspost/compress v1.14.2 // indirect github.com/klauspost/cpuid/v2 v2.0.12 // indirect github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 // indirect github.com/magiconair/properties v1.8.5 // indirect github.com/mattn/go-runewidth v0.0.2 // indirect github.com/matttproud/golang_protobuf_extensions v1.0.1 // indirect github.com/mitchellh/mapstructure v1.5.0 // indirect + github.com/moby/term v0.0.0-20210619224110-3f7ff695adc6 // indirect github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.2 // indirect + github.com/morikuni/aec v1.0.0 // indirect github.com/oklog/ulid v1.3.1 // indirect + github.com/opencontainers/go-digest v1.0.0 // indirect github.com/pierrec/lz4/v4 v4.1.14 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c // indirect github.com/prometheus/common v0.33.0 // indirect github.com/prometheus/procfs v0.7.3 // indirect - github.com/rs/cors v1.8.2 // indirect github.com/sirupsen/logrus v1.7.0 // indirect github.com/soheilhy/cmux v0.1.5 // indirect github.com/spf13/afero v1.6.0 // indirect @@ -162,16 +174,16 @@ require ( go.uber.org/atomic v1.7.0 // indirect go.uber.org/multierr v1.6.0 // indirect go.uber.org/zap v1.17.0 // indirect - golang.org/x/crypto v0.0.0-20220131195533-30dcbda58838 // indirect - golang.org/x/net v0.0.0-20220425223048-2871e0cb64e4 // indirect - golang.org/x/sys v0.0.0-20220310020820-b874c991c1a5 // indirect - golang.org/x/text v0.3.7 // indirect + golang.org/x/crypto v0.1.0 // indirect + golang.org/x/sys v0.1.0 // indirect + golang.org/x/text v0.4.0 // indirect google.golang.org/appengine v1.6.7 // indirect google.golang.org/genproto v0.0.0-20220503193339-ba3ae3f07e29 // indirect gopkg.in/ini.v1 v1.62.0 // indirect gopkg.in/natefinch/lumberjack.v2 v2.0.0 // indirect - gopkg.in/yaml.v3 v3.0.1 // indirect + gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b // indirect + gotest.tools/v3 v3.3.0 // indirect nhooyr.io/websocket v1.8.6 // indirect ) -go 1.19 +go 1.18 diff --git a/go.sum b/go.sum index ddafeb918..df3b1b42e 100644 --- a/go.sum +++ b/go.sum @@ -38,6 +38,8 @@ cloud.google.com/go/storage v1.6.0/go.mod h1:N7U0C8pVQ/+NIKOBQyamJIeKQKkZ+mxpohl cloud.google.com/go/storage v1.8.0/go.mod h1:Wv1Oy7z6Yz3DshWRJFhqM/UCfaWIRTdp0RXyy7KQOVs= cloud.google.com/go/storage v1.10.0/go.mod h1:FLPqc6j+Ki4BU591ie1oL6qBQGu2Bl/tZ9ullr3+Kg0= dmitri.shuralyov.com/gpu/mtl v0.0.0-20190408044501-666a987793e9/go.mod h1:H6x//7gZCb22OMCxBHrMx7a5I7Hp++hsVxbQ4BYO7hU= +github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1 h1:UQHMgLO+TxOElx5B5HZ4hJQsoJ/PvUvKRhJHDQXO8P8= +github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1/go.mod h1:xomTg63KZ2rFqZQzSB4Vz2SUXa1BpHTVz9L5PTmPC4E= github.com/Azure/go-autorest/autorest v0.9.0/go.mod h1:xyHB1BMZT0cuDHU7I0+g046+BFDTQ8rEZB0s4Yfa6bI= github.com/Azure/go-autorest/autorest/adal v0.5.0/go.mod h1:8Z9fGy2MpX0PvDjB1pEgQTmVqjGhiHBW7RJJEciWzS0= github.com/Azure/go-autorest/autorest/date v0.1.0/go.mod h1:plvfp3oPSKwf2DNjlBjWF/7vwR+cUD/ELuzDCXwHUVA= @@ -215,6 +217,14 @@ github.com/dgrijalva/jwt-go v3.2.0+incompatible/go.mod h1:E3ru+11k8xSBh+hMPgOLZm github.com/dgryski/go-farm v0.0.0-20190423205320-6a90982ecee2/go.mod h1:SqUrOPUnsFjfmXRMNPybcSiG0BgUW2AuFH8PAnS2iTw= github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= github.com/dgryski/go-sip13 v0.0.0-20181026042036-e10d5fee7954/go.mod h1:vAd38F8PWV+bWy6jNmig1y/TA+kYO4g3RSRF0IAv0no= +github.com/docker/distribution v2.8.1+incompatible h1:Q50tZOPR6T/hjNsyc9g8/syEs6bk8XXApsHjKukMl68= +github.com/docker/distribution v2.8.1+incompatible/go.mod h1:J2gT2udsDAN96Uj4KfcMRqY0/ypR+oyYUYmja8H+y+w= +github.com/docker/docker v20.10.17+incompatible h1:JYCuMrWaVNophQTOrMMoSwudOVEfcegoZZrleKc1xwE= +github.com/docker/docker v20.10.17+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk= +github.com/docker/go-connections v0.4.0 h1:El9xVISelRB7BuFusrZozjnkIM5YnzCViNKohAFqRJQ= +github.com/docker/go-connections v0.4.0/go.mod h1:Gbd7IOopHjR8Iph03tsViu4nIes5XhDvyHbTtUxmeec= +github.com/docker/go-units v0.4.0 h1:3uh0PgVws3nIA0Q+MwDC8yjEPf9zjRfZZWXZYDct3Tw= +github.com/docker/go-units v0.4.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk= github.com/docker/spdystream v0.0.0-20160310174837-449fdfce4d96/go.mod h1:Qh8CwZgvJUkLughtfhJv5dyTYa91l1fOUCrgjqmcifM= github.com/dustin/go-humanize v0.0.0-20171111073723-bb3d318650d4/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk= github.com/dustin/go-humanize v1.0.0 h1:VSnTsYCnlFHaM2/igO1h6X3HA71jcobQuxemgkq4zYo= @@ -343,12 +353,15 @@ github.com/gobuffalo/packd v0.1.0/go.mod h1:M2Juc+hhDXf/PnmBANFCqx4DM3wRbgDvnVWe github.com/gobuffalo/packr/v2 v2.0.9/go.mod h1:emmyGweYTm6Kdper+iywB6YK5YzuKchGtJQZ0Odn4pQ= github.com/gobuffalo/packr/v2 v2.2.0/go.mod h1:CaAwI0GPIAv+5wKLtv8Afwl+Cm78K/I/VCm/3ptBN+0= github.com/gobuffalo/syncx v0.0.0-20190224160051-33c29581e754/go.mod h1:HhnNqWY95UYwwW3uSASeV7vtgYkT2t16hJgV3AEPUpw= -github.com/gobwas/httphead v0.0.0-20180130184737-2c6c146eadee h1:s+21KNqlpePfkah2I+gwHF8xmJWRjooY+5248k6m4A0= github.com/gobwas/httphead v0.0.0-20180130184737-2c6c146eadee/go.mod h1:L0fX3K22YWvt/FAX9NnzrNzcI4wNYi9Yku4O0LKYflo= -github.com/gobwas/pool v0.2.0 h1:QEmUOlnSjWtnpRGHF3SauEiOsy82Cup83Vf2LcMlnc8= +github.com/gobwas/httphead v0.0.0-20200921212729-da3d93bc3c58 h1:YyrUZvJaU8Q0QsoVo+xLFBgWDTam29PKea6GYmwvSiQ= +github.com/gobwas/httphead v0.0.0-20200921212729-da3d93bc3c58/go.mod h1:L0fX3K22YWvt/FAX9NnzrNzcI4wNYi9Yku4O0LKYflo= github.com/gobwas/pool v0.2.0/go.mod h1:q8bcK0KcYlCgd9e7WYLm9LpyS+YeLd8JVDW6WezmKEw= -github.com/gobwas/ws v1.0.2 h1:CoAavW/wd/kulfZmSIBt6p24n4j7tHgNVCjsfHVNUbo= +github.com/gobwas/pool v0.2.1 h1:xfeeEhW7pwmX8nuLVlqbzVc7udMDrwetjEv+TZIz1og= +github.com/gobwas/pool v0.2.1/go.mod h1:q8bcK0KcYlCgd9e7WYLm9LpyS+YeLd8JVDW6WezmKEw= github.com/gobwas/ws v1.0.2/go.mod h1:szmBTxLgaFppYjEmNtny/v3w89xOydFnnZMcgRRu/EM= +github.com/gobwas/ws v1.0.4 h1:5eXU1CZhpQdq5kXbKb+sECH5Ia5KiO6CYzIzdlVx6Bs= +github.com/gobwas/ws v1.0.4/go.mod h1:szmBTxLgaFppYjEmNtny/v3w89xOydFnnZMcgRRu/EM= github.com/gocql/gocql v0.0.0-20220224095938-0eacd3183625/go.mod h1:3gM2c4D3AnkISwBxGnMMsS8Oy4y2lhbPRsH4xnJrHG8= github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA= github.com/gofiber/fiber/v2 v2.11.0/go.mod h1:oZTLWqYnqpMMuF922SjGbsYZsdpE1MCfh416HNdweIM= @@ -615,8 +628,8 @@ github.com/jackc/puddle v1.1.0/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dv github.com/jackc/puddle v1.1.1/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk= github.com/jackc/puddle v1.1.3/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk= github.com/jackc/puddle v1.2.0/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk= -github.com/jaffee/commandeer v0.5.0 h1:241M9N+gHQmPyjIG+yy8GGcZPfzFuIyOmJHzm5ka92g= -github.com/jaffee/commandeer v0.5.0/go.mod h1:kCwfuSvZ2T0NVEr3LDSo6fDUgi0xSBnAVDdkOKTtpLQ= +github.com/jaffee/commandeer v0.6.0 h1:YI44XLWcJN21euhh32sZW8vM/tljPYxhsXIfEPkQKcs= +github.com/jaffee/commandeer v0.6.0/go.mod h1:kCwfuSvZ2T0NVEr3LDSo6fDUgi0xSBnAVDdkOKTtpLQ= github.com/jedib0t/go-pretty v4.3.0+incompatible h1:CGs8AVhEKg/n9YbUenWmNStRW2PHJzaeDodcfvRAbIo= github.com/jedib0t/go-pretty v4.3.0+incompatible/go.mod h1:XemHduiw8R651AF9Pt4FwCTKeG3oo7hrHJAoznj9nag= github.com/jhump/gopoet v0.0.0-20190322174617-17282ff210b3/go.mod h1:me9yfT6IJSlOL3FCfrg+L6yzUEZ+5jW6WHt4Sk+UPUI= @@ -636,9 +649,8 @@ github.com/jmespath/go-jmespath/internal/testify v1.5.1/go.mod h1:L3OGu8Wl2/fWfC github.com/jmoiron/sqlx v1.2.0/go.mod h1:1FEQNm3xlJgrMD+FBdI9+xvCksHtbpVBBw5dYhBSsks= github.com/joho/godotenv v1.3.0/go.mod h1:7hK45KPybAkOC6peb+G5yklZfMxEjkZhHbwpqxOKXbg= github.com/jonboulle/clockwork v0.1.0/go.mod h1:Ii8DK3G1RaLaWxj9trq07+26W01tbo22gdxWY5EU2bo= +github.com/jonboulle/clockwork v0.2.2 h1:UOGuzwb1PwsrDAObMuhUnj0p5ULPj8V/xJ7Kx9qUBdQ= github.com/jonboulle/clockwork v0.2.2/go.mod h1:Pkfl5aHPm1nk2H9h0bjmnJD/BcgbGXUBGnn1kMkgxc8= -github.com/jonboulle/clockwork v0.3.0 h1:9BSCMi8C+0qdApAp4auwX0RkLGUjs956h0EkuQymUhg= -github.com/jonboulle/clockwork v0.3.0/go.mod h1:Pkfl5aHPm1nk2H9h0bjmnJD/BcgbGXUBGnn1kMkgxc8= github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY= github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y= github.com/jpillora/backoff v1.0.0/go.mod h1:J/6gKK9jxlEcS3zixgDgUAsiuZ7yrSoa/FX5e0EB2j4= @@ -672,9 +684,8 @@ github.com/klauspost/compress v1.11.7/go.mod h1:aoV0uJVorq1K+umq18yTdKaF57EivdYs github.com/klauspost/compress v1.12.2/go.mod h1:8dP1Hq4DHOhN9w426knH3Rhby4rFm6D8eO+e+Dq5Gzg= github.com/klauspost/compress v1.13.4/go.mod h1:8dP1Hq4DHOhN9w426knH3Rhby4rFm6D8eO+e+Dq5Gzg= github.com/klauspost/compress v1.13.6/go.mod h1:/3/Vjq9QcHkK5uEr5lBEmyoZ1iFhe47etQ6QUkpK6sk= +github.com/klauspost/compress v1.14.2 h1:S0OHlFk/Gbon/yauFJ4FfJJF5V0fc5HbBTJazi28pRw= github.com/klauspost/compress v1.14.2/go.mod h1:/3/Vjq9QcHkK5uEr5lBEmyoZ1iFhe47etQ6QUkpK6sk= -github.com/klauspost/compress v1.15.1 h1:y9FcTHGyrebwfP0ZZqFiaxTaiDnUrGkJkI+f583BL1A= -github.com/klauspost/compress v1.15.1/go.mod h1:/3/Vjq9QcHkK5uEr5lBEmyoZ1iFhe47etQ6QUkpK6sk= github.com/klauspost/cpuid/v2 v2.0.12 h1:p9dKCg8i4gmOxtv35DvrYoWqYzQrvEVdjQ762Y0OqZE= github.com/klauspost/cpuid/v2 v2.0.12/go.mod h1:g2LTdtYhdyuGPqyWyv7qRAmj1WBqxuObKfj5c0PQa7c= github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ= @@ -768,6 +779,8 @@ github.com/mitchellh/mapstructure v1.4.2/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RR github.com/mitchellh/mapstructure v1.5.0 h1:jeMsZIYE/09sWLaz43PL7Gy6RuMjD2eJVyuac5Z2hdY= github.com/mitchellh/mapstructure v1.5.0/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RRV2QTWOzhPopBRo= github.com/mitchellh/reflectwalk v1.0.0/go.mod h1:mSTlrgnPZtwu0c4WaC2kGObEpuNDbx0jmZXqmk4esnw= +github.com/moby/term v0.0.0-20210619224110-3f7ff695adc6 h1:dcztxKSvZ4Id8iPpHERQBbIJfabdt4wUm5qy3wOL2Zc= +github.com/moby/term v0.0.0-20210619224110-3f7ff695adc6/go.mod h1:E2VnQOmVuvZB6UYnnDB0qG5Nq/1tD9acaOpo6xmt0Kw= github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg= github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= @@ -779,6 +792,8 @@ github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjY github.com/molecula/apophenia v0.0.0-20190827192002-68b7a14a478b h1:cZADDaNYM7xn/nklO3g198JerGQjadFuA0ofxBJgK0Y= github.com/molecula/apophenia v0.0.0-20190827192002-68b7a14a478b/go.mod h1:uXd1BiH7xLmgkhVmspdJLENv6uGWrTL/MQX2TN7Yz9s= github.com/montanaflynn/stats v0.0.0-20171201202039-1bf9dbcd8cbe/go.mod h1:wL8QJuTMNUDYhXwkmfOly8iTdp5TEcJFWZD2D7SIkUc= +github.com/morikuni/aec v1.0.0 h1:nP9CBfwrvYnBRgY6qfDQkygYDmYwOilePFkwzv4dU8A= +github.com/morikuni/aec v1.0.0/go.mod h1:BbKIizmSmc5MMPqRYbxO4ZU0S0+P200+tUnFx7PXmsc= github.com/munnerz/goautoneg v0.0.0-20120707110453-a547fc61f48d/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ= github.com/mwitkow/go-conntrack v0.0.0-20161129095857-cc309e4a2223/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U= github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f h1:KUppIJq7/+SVif2QVs3tOP0zanoHgBEVAwHxUSIzRqU= @@ -818,6 +833,10 @@ github.com/onsi/gomega v1.10.1/go.mod h1:iN09h71vgCQne3DLsj+A5owkum+a2tYe+TOCB1y github.com/onsi/gomega v1.10.2/go.mod h1:iN09h71vgCQne3DLsj+A5owkum+a2tYe+TOCB1ybHNo= github.com/onsi/gomega v1.16.0/go.mod h1:HnhC7FXeEQY45zxNK3PPoIUhzk/80Xly9PcubAlGdZY= github.com/op/go-logging v0.0.0-20160315200505-970db520ece7/go.mod h1:HzydrMdWErDVzsI23lYNej1Htcns9BCg93Dk0bBINWk= +github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U= +github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM= +github.com/opencontainers/image-spec v1.0.2 h1:9yCKha/T5XdGtO0q9Q9a6T5NUCsTn/DrBg0D7ufOcFM= +github.com/opencontainers/image-spec v1.0.2/go.mod h1:BtxoFyWECRxE4U/7sNtV5W15zMzWCbyJoFRP3s7yZA0= github.com/opentracing-contrib/go-observer v0.0.0-20170622124052-a52f23424492/go.mod h1:Ngi6UdF0k5OKD5t5wlmGhe/EDKPoUM3BXZSSfIuJbis= github.com/opentracing/basictracer-go v1.0.0/go.mod h1:QfBfYuafItcjQuMwinw9GhYKwFXS9KnPs5lxoYwgW74= github.com/opentracing/opentracing-go v1.0.2/go.mod h1:UkNAQd3GIcIGf0SeVgPpRdFStlNbqXla1AfSYxPUl2o= @@ -938,7 +957,6 @@ github.com/santhosh-tekuri/jsonschema/v5 v5.0.0/go.mod h1:FKdcjfQW6rpZSnxxUvEA5H github.com/satori/go.uuid v1.2.0/go.mod h1:dA0hQrYB0VpLJoorglMZABFdXlWrHn1NEOzdhQKdks0= github.com/satori/go.uuid v1.2.1-0.20180404165556-75cca531ea76 h1:ofyVTM1w4iyKwaQIlRR6Ip06mXXx5Cnz7a4mTGYq1hE= github.com/satori/go.uuid v1.2.1-0.20180404165556-75cca531ea76/go.mod h1:dA0hQrYB0VpLJoorglMZABFdXlWrHn1NEOzdhQKdks0= -github.com/schollz/closestmatch v2.1.0+incompatible/go.mod h1:RtP1ddjLong6gTkbtmuhtR2uUrrJOpYzYRvbcPAid+g= github.com/sean-/seed v0.0.0-20170313163322-e2103e2c3529/go.mod h1:DxrIzT+xaE7yg65j358z/aeFdxmN0P9QXhEzd20vsDc= github.com/segmentio/kafka-go v0.4.29 h1:4ujULpikzHG0HqKhjumDghFjy/0RRCSl/7lbriwQAH0= github.com/segmentio/kafka-go v0.4.29/go.mod h1:m1lXeqJtIFYZayv0shM/tjrAFljvWLTprxBHd+3PnaU= @@ -1198,8 +1216,9 @@ golang.org/x/crypto v0.0.0-20210513164829-c07d793c2f9a/go.mod h1:P+XmwS30IXTQdn5 golang.org/x/crypto v0.0.0-20210616213533-5ff15b29337e/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc= golang.org/x/crypto v0.0.0-20210711020723-a769d52b0f97/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc= golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc= -golang.org/x/crypto v0.0.0-20220131195533-30dcbda58838 h1:71vQrMauZZhcTVK6KdYM+rklehEEwb3E+ZhaE5jrPrE= golang.org/x/crypto v0.0.0-20220131195533-30dcbda58838/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4= +golang.org/x/crypto v0.1.0 h1:MDRAIl0xIo9Io2xV565hzXHw3zVseKrJKodhohM5CjU= +golang.org/x/crypto v0.1.0/go.mod h1:RecgLatLF4+eUMCP1PoPZQb+cVrJcOPbHkTkbkB9sbw= golang.org/x/exp v0.0.0-20180321215751-8460e604b9de/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= golang.org/x/exp v0.0.0-20180807140117-3d87b88a115f/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= @@ -1216,8 +1235,8 @@ golang.org/x/exp v0.0.0-20200224162631-6cc2880d07d6/go.mod h1:3jZMyOhIsHpP37uCMk golang.org/x/exp v0.0.0-20200331195152-e8c3332aa8e5/go.mod h1:4M0jN8W1tt0AVLNr8HDosyJCDCDuyL9N9+3m7wDWgKw= golang.org/x/exp v0.0.0-20200901203048-c4f52b2c50aa/go.mod h1:3jZMyOhIsHpP37uCMkUooju7aAi5cS1Q23tOzKc+0MU= golang.org/x/exp v0.0.0-20200908183739-ae8ad444f925/go.mod h1:1phAWC201xIgDyaFpmDeZkgf70Q4Pd/CNqfRtVPtxNw= -golang.org/x/exp v0.0.0-20201008143054-e3b2a7f2fdc7 h1:2/QncOxxpPAdiH+E00abYw/SaQG353gltz79Nl1zrYE= -golang.org/x/exp v0.0.0-20201008143054-e3b2a7f2fdc7/go.mod h1:1phAWC201xIgDyaFpmDeZkgf70Q4Pd/CNqfRtVPtxNw= +golang.org/x/exp v0.0.0-20221031165847-c99f073a8326 h1:QfTh0HpN6hlw6D3vu8DAwC8pBIwikq0AI1evdm+FksE= +golang.org/x/exp v0.0.0-20221031165847-c99f073a8326/go.mod h1:CxIveKay+FTh1D0yPZemJVgC/95VzuuOLq5Qi4xnoYc= golang.org/x/image v0.0.0-20180708004352-c73c2afc3b81/go.mod h1:ux5Hcp/YLpHSI86hEcLt0YII63i6oz57MZXIpbrjZUs= golang.org/x/image v0.0.0-20190227222117-0694c2d4d067/go.mod h1:kZ7UVZpmo3dzQBMxlp+ypCbDeSB+sBbTgSJuh5dn5js= golang.org/x/image v0.0.0-20190802002840-cff245a6509b/go.mod h1:FeLwcggjj3mMvU+oOTbSwawSJRM1uh48EjtB4UJZlP0= @@ -1245,8 +1264,8 @@ golang.org/x/mod v0.3.1-0.20200828183125-ce943fd02449/go.mod h1:s0Qsj1ACt9ePp/hM golang.org/x/mod v0.4.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= golang.org/x/mod v0.4.1/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= golang.org/x/mod v0.4.2/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= -golang.org/x/mod v0.5.1 h1:OJxoQ/rynoF0dcCdI7cLPktw/hR2cueqYfjm43oqK38= -golang.org/x/mod v0.5.1/go.mod h1:5OXOZSfqPIIbmVBIIKWRFfZjPR0E5r58TLhUjH0a2Ro= +golang.org/x/mod v0.6.0 h1:b9gGHsz9/HhJ3HF5DHQytPpuwocVTChQJK3AvoLRD5I= +golang.org/x/mod v0.6.0/go.mod h1:4mET923SAdbXp2ki8ey+zGs1SLqsuM2Y0uvdZR/fUNI= golang.org/x/net v0.0.0-20170114055629-f2499483f923/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= @@ -1308,8 +1327,8 @@ golang.org/x/net v0.0.0-20211112202133-69e39bad7dc2/go.mod h1:9nx3DQGgdP8bBQD5qx golang.org/x/net v0.0.0-20211216030914-fe4d6282115f/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y= golang.org/x/net v0.0.0-20220127200216-cd36cc0744dd/go.mod h1:CfG3xpIq0wQ8r1q4Su4UZFWDARRcnwPjda9FqA0JpMk= golang.org/x/net v0.0.0-20220225172249-27dd8689420f/go.mod h1:CfG3xpIq0wQ8r1q4Su4UZFWDARRcnwPjda9FqA0JpMk= -golang.org/x/net v0.0.0-20220425223048-2871e0cb64e4 h1:HVyaeDAYux4pnY+D/SiwmLOR36ewZ4iGQIIrtnuCjFA= -golang.org/x/net v0.0.0-20220425223048-2871e0cb64e4/go.mod h1:CfG3xpIq0wQ8r1q4Su4UZFWDARRcnwPjda9FqA0JpMk= +golang.org/x/net v0.1.0 h1:hZ/3BUoy5aId7sCpA/Tc5lt8DkFgdVS2onTpJsZ/fl0= +golang.org/x/net v0.1.0/go.mod h1:Cx3nUiGt4eDBEyega/BKRp+/AlGL8hYe7U9odMt2Cco= golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U= golang.org/x/oauth2 v0.0.0-20190226205417-e64efc72b421/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw= golang.org/x/oauth2 v0.0.0-20190604053449-0f29369cfe45/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw= @@ -1419,6 +1438,7 @@ golang.org/x/sys v0.0.0-20210510120138-977fb7262007/go.mod h1:oPkhp1MJrh7nUepCBc golang.org/x/sys v0.0.0-20210514084401-e8d321eab015/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20210603081109-ebe580a85c40/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.0.0-20210616094352-59db8d763f22/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20210809222454-d867a43fc93e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20210927094055-39ccf1dd6fa6/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= @@ -1428,8 +1448,9 @@ golang.org/x/sys v0.0.0-20211216021012-1d35b9e2eb4e/go.mod h1:oPkhp1MJrh7nUepCBc golang.org/x/sys v0.0.0-20220114195835-da31bd327af9/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220128215802-99c3d69c2c27/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220227234510-4e6760a101f9/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.0.0-20220310020820-b874c991c1a5 h1:y/woIyUBFbpQGKS0u1aHF/40WUDnek3fPOyD08H5Vng= golang.org/x/sys v0.0.0-20220310020820-b874c991c1a5/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.1.0 h1:kunALQeHf1/185U1i0GOB/fy1IPRDDpuoOOqRReG57U= +golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/term v0.0.0-20201117132131-f5c789dd3221/go.mod h1:Nr5EML6q2oocZ2LXRh80K7BxOlk5/8JxuGnuhpl+muw= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8= @@ -1442,8 +1463,9 @@ golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.3.4/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.3.5/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= -golang.org/x/text v0.3.7 h1:olpwvP2KacW1ZWvsR7uQhoyTYvKAupfQrRGBFM352Gk= golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ= +golang.org/x/text v0.4.0 h1:BrVqGRd7+k1DiOgtnFvAkoQEWQvBc25ouMJM6429SFg= +golang.org/x/text v0.4.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= golang.org/x/time v0.0.0-20180412165947-fbb02b2291d2/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= @@ -1474,6 +1496,7 @@ golang.org/x/tools v0.0.0-20190524140312-2c0ae7006135/go.mod h1:RgjU9mgBXZiqYHBn golang.org/x/tools v0.0.0-20190531172133-b3315ee88b7d/go.mod h1:/rFqwRUd4F7ZHNgwSSTFct+R/Kf4OFW1sUzUTQQTgfc= golang.org/x/tools v0.0.0-20190606124116-d0a3d012864b/go.mod h1:/rFqwRUd4F7ZHNgwSSTFct+R/Kf4OFW1sUzUTQQTgfc= golang.org/x/tools v0.0.0-20190621195816-6e04913cbbac/go.mod h1:/rFqwRUd4F7ZHNgwSSTFct+R/Kf4OFW1sUzUTQQTgfc= +golang.org/x/tools v0.0.0-20190624222133-a101b041ded4/go.mod h1:/rFqwRUd4F7ZHNgwSSTFct+R/Kf4OFW1sUzUTQQTgfc= golang.org/x/tools v0.0.0-20190628153133-6cdbf07be9d0/go.mod h1:/rFqwRUd4F7ZHNgwSSTFct+R/Kf4OFW1sUzUTQQTgfc= golang.org/x/tools v0.0.0-20190816200558-6889da9d5479/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.0.0-20190823170909-c4a336ef6a2f/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= @@ -1522,8 +1545,8 @@ golang.org/x/tools v0.0.0-20210105154028-b0ab187a4818/go.mod h1:emZCQorbCU4vsT4f golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= golang.org/x/tools v0.1.0/go.mod h1:xkSsbof2nBLbhDlRMhhhyNLN/zl3eTqcnHD5viDpcZ0= golang.org/x/tools v0.1.1/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk= -golang.org/x/tools v0.1.2 h1:kRBLX7v7Af8W7Gdbbc908OJcdgtK8bOz9Uaj8/F1ACA= golang.org/x/tools v0.1.2/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk= +golang.org/x/tools v0.2.0 h1:G6AHpWxTMGY1KyEYoAQ5WTtIekUUvDNjan3ugu60JvE= golang.org/x/xerrors v0.0.0-20190410155217-1f06c39b4373/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20190513163551-3ee3066db522/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= @@ -1687,6 +1710,9 @@ gorm.io/driver/sqlserver v1.0.4/go.mod h1:ciEo5btfITTBCj9BkoUVDvgQbUdLWQNqdFY5OG gorm.io/gorm v1.9.19/go.mod h1:0HFTzE/SqkGTzK6TlDPPQbAYCluiVvhzoA1+aVyzenw= gorm.io/gorm v1.20.0/go.mod h1:0HFTzE/SqkGTzK6TlDPPQbAYCluiVvhzoA1+aVyzenw= gorm.io/gorm v1.20.6/go.mod h1:0HFTzE/SqkGTzK6TlDPPQbAYCluiVvhzoA1+aVyzenw= +gotest.tools/v3 v3.0.2/go.mod h1:3SzNCllyD9/Y+b5r9JIKQ474KzkZyqLqEfYqMsX94Bk= +gotest.tools/v3 v3.3.0 h1:MfDY1b1/0xN1CyMlQDac0ziEy9zJQd9CXBRRDHw2jJo= +gotest.tools/v3 v3.3.0/go.mod h1:Mcr9QNxkg0uMvy/YElmo4SpXgJKWgQvYrT7Kw5RzJ1A= honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= honnef.co/go/tools v0.0.0-20190106161140-3f1c8253044a/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= honnef.co/go/tools v0.0.0-20190418001031-e561f6794a2a/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= diff --git a/handler.go b/handler.go index 3cdaffea2..f0a6a0a78 100644 --- a/handler.go +++ b/handler.go @@ -427,6 +427,7 @@ type ImportRoaringRequest struct { Block int Views map[string][]byte UpdateExistence bool + SuppressLog bool } // ImportRoaringShardRequest is the request for the shard @@ -438,6 +439,11 @@ type ImportRoaringShardRequest struct { // a successful response to the client. Remote bool Views []RoaringUpdate + + // SuppressLog requests we not write to the write log. Typically + // that would be because this request is being replayed from a + // write log. + SuppressLog bool } // RoaringUpdate represents the bits to clear and then set in a particular view. diff --git a/holder.go b/holder.go index 4b5d270be..55f356a54 100644 --- a/holder.go +++ b/holder.go @@ -13,14 +13,15 @@ import ( "sync" "time" - "github.com/featurebasedb/featurebase/v3/disco" - "github.com/featurebasedb/featurebase/v3/logger" - rbfcfg "github.com/featurebasedb/featurebase/v3/rbf/cfg" - "github.com/featurebasedb/featurebase/v3/roaring" - "github.com/featurebasedb/featurebase/v3/stats" - "github.com/featurebasedb/featurebase/v3/storage" - "github.com/featurebasedb/featurebase/v3/testhook" - "github.com/featurebasedb/featurebase/v3/vprint" + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/disco" + "github.com/molecula/featurebase/v3/logger" + rbfcfg "github.com/molecula/featurebase/v3/rbf/cfg" + "github.com/molecula/featurebase/v3/roaring" + "github.com/molecula/featurebase/v3/stats" + "github.com/molecula/featurebase/v3/storage" + "github.com/molecula/featurebase/v3/testhook" + "github.com/molecula/featurebase/v3/vprint" "github.com/pkg/errors" "golang.org/x/sync/errgroup" ) @@ -132,6 +133,11 @@ type Holder struct { // on holding mu. imu sync.RWMutex indexes map[string]*Index + + // directive is the latest directive applied to the node. + directive *dax.Directive + + versionStore dax.VersionStore } // HolderOpts holds information about the holder which other things might want @@ -142,6 +148,31 @@ type HolderOpts struct { StorageBackend string } +func (h *Holder) Directive() dax.Directive { + h.mu.RLock() + defer h.mu.RUnlock() + + if h.directive == nil { + return dax.Directive{} + } + return *h.directive +} + +func (h *Holder) SetDirective(d *dax.Directive) { + if d == nil { + return + } + + h.mu.Lock() + defer h.mu.Unlock() + + // Only set the cached directive if the incoming version is newer than that + // of the existing directive's version. + if h.directive == nil || d.Version > h.directive.Version { + h.directive = d + } +} + func (h *Holder) StartTransaction(ctx context.Context, id string, timeout time.Duration, exclusive bool) (*Transaction, error) { return h.transactionManager.Start(ctx, id, timeout, exclusive) } @@ -284,6 +315,8 @@ func NewHolder(path string, cfg *HolderConfig) *Holder { path: path, indexes: make(map[string]*Index), + + versionStore: dax.NewNopVersionStore(), } txf, err := NewTxFactory(cfg.StorageConfig.Backend, h.IndexesPath(), h) @@ -957,6 +990,14 @@ func (h *Holder) createIndex(cim *CreateIndexMessage, broadcast bool) (*Index, e // Update options. h.addIndex(index) + tkey := dax.TableKey(cim.Index) + qtid := tkey.QualifiedTableID() + + // Initialize the table in holder.versionStore. + if err := h.versionStore.AddTable(context.Background(), qtid); err != nil { + h.Logger.Printf("could not add table to version store: %s", cim.Index) + } + if broadcast { // Send the create index message to all nodes. if err := h.broadcaster.SendSync(cim); err != nil { @@ -973,6 +1014,59 @@ func (h *Holder) createIndex(cim *CreateIndexMessage, broadcast bool) (*Index, e return index, nil } +// createIndexWithPartitions is similar to createIndex, but it takes a list of +// partitions for which this node is responsible. This ensures that the node +// doesn't instantiate more partition TranslateStores than is necessary. +func (h *Holder) createIndexWithPartitions(cim *CreateIndexMessage, translatePartitions dax.Partitions) (*Index, error) { + if cim.Index == "" { + return nil, errors.New("index name required") + } + + // Otherwise create a new index. + index, err := h.newIndex(h.IndexPath(cim.Index), cim.Index) + if err != nil { + return nil, errors.Wrap(err, "creating") + } + + index.keys = cim.Meta.Keys + index.trackExistence = cim.Meta.TrackExistence + index.createdAt = cim.CreatedAt + index.translatePartitions = translatePartitions + + if err = index.Open(); err != nil { + return nil, errors.Wrap(err, "opening") + } + + // Update options. + h.addIndex(index) + + tkey := dax.TableKey(cim.Index) + qtid := tkey.QualifiedTableID() + + // Initialize the table in holder.versionStore. + if err := h.versionStore.AddTable(context.Background(), qtid); err != nil { + h.Logger.Printf("could not add table to version store: %s", cim.Index) + } + + // Initialize a list of partitions at version 0. + newPartitions := make(dax.Partitions, len(translatePartitions)) + for i := range translatePartitions { + newPartitions[i] = dax.NewPartition(translatePartitions[i].Num, 0) + } + + if err := h.versionStore.AddPartitions(context.Background(), qtid, newPartitions...); err != nil { + return nil, errors.Wrap(err, "adding partitions to version store") + } + + // Since this is a new index, we need to kick off + // its translation sync. + if err := h.translationSyncer.Reset(); err != nil { + return nil, errors.Wrap(err, "resetting translation sync") + } + + return index, nil +} + func (h *Holder) loadSchema() error { schema, err := h.Schemator.Schema(context.TODO()) if err != nil { @@ -1077,7 +1171,11 @@ func (h *Holder) newIndex(path, name string) (*Index, error) { func (h *Holder) DeleteIndex(name string) error { h.mu.Lock() defer h.mu.Unlock() + return h.deleteIndex(name) +} +// deleteIndex is a non-locking version of DeleteIndex(). +func (h *Holder) deleteIndex(name string) error { // Confirm index exists. index := h.Index(name) if index == nil { @@ -1114,7 +1212,15 @@ func (h *Holder) DeleteIndex(name string) error { } // Remove reference. - h.deleteIndex(name) + h.deleteIndexFromMap(name) + + tkey := dax.TableKey(name) + qtid := tkey.QualifiedTableID() + + // Remove the index from holder.versionStore. + if _, _, err := h.versionStore.RemoveTable(context.Background(), qtid); err != nil { + h.Logger.Printf("could not find table to remove from version store: %s", name) + } // I'm not sure if calling Reset() here is necessary // since closing the index stops its translation @@ -1122,7 +1228,7 @@ func (h *Holder) DeleteIndex(name string) error { return h.translationSyncer.Reset() } -func (h *Holder) deleteIndex(index string) { +func (h *Holder) deleteIndexFromMap(index string) { h.imu.Lock() delete(h.indexes, index) h.imu.Unlock() diff --git a/http_handler.go b/http_handler.go index 94a1e293e..c3520b4eb 100644 --- a/http_handler.go +++ b/http_handler.go @@ -41,6 +41,17 @@ import ( "github.com/felixge/fgprof" "github.com/gorilla/handlers" "github.com/gorilla/mux" + "github.com/molecula/featurebase/v3/authn" + "github.com/molecula/featurebase/v3/authz" + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/disco" + "github.com/molecula/featurebase/v3/logger" + "github.com/molecula/featurebase/v3/monitor" + "github.com/molecula/featurebase/v3/pql" + "github.com/molecula/featurebase/v3/rbf" + "github.com/molecula/featurebase/v3/sql3/planner/types" + "github.com/molecula/featurebase/v3/storage" + "github.com/molecula/featurebase/v3/tracing" "github.com/pkg/errors" "github.com/prometheus/client_golang/prometheus/promhttp" dto "github.com/prometheus/client_model/go" @@ -310,7 +321,7 @@ func (h *Handler) populateValidators() { h.validators["PostImport"] = queryValidationSpecRequired().Optional("clear", "ignoreKeyCheck") h.validators["PostImportAtomicRecord"] = queryValidationSpecRequired().Optional("simPowerLossAfter") h.validators["PostImportRoaring"] = queryValidationSpecRequired().Optional("remote", "clear") - h.validators["PostQuery"] = queryValidationSpecRequired().Optional("shards", "excludeColumns", "profile") + h.validators["PostQuery"] = queryValidationSpecRequired().Optional("shards", "excludeColumns", "profile", "remote") h.validators["GetInfo"] = queryValidationSpecRequired() h.validators["RecalculateCaches"] = queryValidationSpecRequired() h.validators["GetSchema"] = queryValidationSpecRequired().Optional("views") @@ -609,6 +620,12 @@ func newRouter(handler *Handler) http.Handler { router.HandleFunc("/userinfo", handler.handleUserInfo).Methods("GET").Name("UserInfo") router.HandleFunc("/internal/oauth-config", handler.handleOAuthConfig).Methods("GET").Name("GetOAuthConfig") + router.HandleFunc("/health", handler.handleGetHealth).Methods("GET").Name("GetHealth") + router.HandleFunc("/directive", handler.handlePostDirective).Methods("POST").Name("PostDirective") + router.HandleFunc("/snapshot/shard-data", handler.handlePostSnapshotShardData).Methods("POST").Name("PostShapshotShardData") + router.HandleFunc("/snapshot/table-keys", handler.handlePostSnapshotTableKeys).Methods("POST").Name("PostShapshotTableKeys") + router.HandleFunc("/snapshot/field-keys", handler.handlePostSnapshotFieldKeys).Methods("POST").Name("PostShapshotFieldKeys") + // Endpoints to support lattice UI embedded via statik. // The messiness here reflects the fact that assets live in a nontrivial // directory structure that is controlled externally. @@ -2472,6 +2489,8 @@ func (h *Handler) readURLQueryRequest(r *http.Request) (*QueryRequest, error) { return nil, errors.New("invalid shard argument") } + remote := parseBool(q.Get("remote")) + // Optional profiling profile := false profileString := q.Get("profile") @@ -2484,11 +2503,16 @@ func (h *Handler) readURLQueryRequest(r *http.Request) (*QueryRequest, error) { return &QueryRequest{ Query: query, + Remote: remote, Shards: shards, Profile: profile, }, nil } +func parseBool(a string) bool { + return strings.ToLower(a) == "true" +} + // writeQueryResponse writes the response from the executor to w. func (h *Handler) writeQueryResponse(w http.ResponseWriter, r *http.Request, resp *QueryResponse) error { if !validHeaderAcceptJSON(r.Header) { @@ -3544,7 +3568,6 @@ func (h *Handler) handleFindOrCreateKeys(w http.ResponseWriter, r *http.Request, translations, err = h.api.CreateIndexKeys(r.Context(), indexName, keys...) case !requireField && !create: translations, err = h.api.FindIndexKeys(r.Context(), indexName, keys...) - } if err != nil { http.Error(w, fmt.Sprintf("translating keys: %v", err), http.StatusInternalServerError) @@ -3896,3 +3919,103 @@ func getTokens(r *http.Request) (string, string) { return access, refresh } + +func (h *Handler) handlePostDirective(w http.ResponseWriter, r *http.Request) { + if !validHeaderAcceptJSON(r.Header) { + http.Error(w, "JSON only acceptable response", http.StatusNotAcceptable) + return + } + + body := r.Body + defer body.Close() + + d := &dax.Directive{} + if err := json.NewDecoder(body).Decode(d); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + if err := h.api.Directive(r.Context(), d); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + w.WriteHeader(http.StatusOK) +} + +// POST /snapshot/shard-data +func (h *Handler) handlePostSnapshotShardData(w http.ResponseWriter, r *http.Request) { + if !validHeaderAcceptJSON(r.Header) { + http.Error(w, "JSON only acceptable response", http.StatusNotAcceptable) + return + } + + body := r.Body + defer body.Close() + + req := &dax.SnapshotShardDataRequest{} + if err := json.NewDecoder(body).Decode(req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + if err := h.api.SnapshotShardData(r.Context(), req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + w.WriteHeader(http.StatusOK) +} + +// POST /snapshot/table-keys +func (h *Handler) handlePostSnapshotTableKeys(w http.ResponseWriter, r *http.Request) { + if !validHeaderAcceptJSON(r.Header) { + http.Error(w, "JSON only acceptable response", http.StatusNotAcceptable) + return + } + + body := r.Body + defer body.Close() + + req := &dax.SnapshotTableKeysRequest{} + if err := json.NewDecoder(body).Decode(req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + if err := h.api.SnapshotTableKeys(r.Context(), req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + w.WriteHeader(http.StatusOK) +} + +// POST /snapshot/field-keys +func (h *Handler) handlePostSnapshotFieldKeys(w http.ResponseWriter, r *http.Request) { + if !validHeaderAcceptJSON(r.Header) { + http.Error(w, "JSON only acceptable response", http.StatusNotAcceptable) + return + } + + body := r.Body + defer body.Close() + + req := &dax.SnapshotFieldKeysRequest{} + if err := json.NewDecoder(body).Decode(req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + if err := h.api.SnapshotFieldKeys(r.Context(), req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + + w.WriteHeader(http.StatusOK) +} + +// GET /health +func (h *Handler) handleGetHealth(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) +} diff --git a/idk/Makefile b/idk/Makefile index 80bae7b61..23c3c4172 100644 --- a/idk/Makefile +++ b/idk/Makefile @@ -9,6 +9,7 @@ BINOUT ?= bin VERSION := $(shell git describe --tags 2> /dev/null || git rev-parse --verify --short=7 HEAD) VERSION_ID = $(VERSION)-$(GOOS)-$(GOARCH) BUILD_TIME := $(shell date -u +%FT%T%z) +AWS_ACCOUNTID ?= undefined BRANCH_NAME ?= "" # We allow setting a custom docker-compose "project". Multiple of the @@ -141,10 +142,11 @@ save-%-logs: $(DOCKER_COMPOSE) logs $* > ./testdata/$(PROJECT)_$*_logs.txt -start-all: .pulled testenv build-wait +start-all: testenv build-wait echo "branch name" ${BRANCH_NAME} $(DOCKER_COMPOSE) up -d zookeeper $(DOCKER_COMPOSE) run -T wait zookeeper 'echo "ruok" | nc -w 2 zookeeper 2181 | grep imok' + $(DOCKER_COMPOSE) up -d dax $(DOCKER_COMPOSE) up -d kafka $(DOCKER_COMPOSE) up -d schema-registry $(DOCKER_COMPOSE) up -d postgres @@ -155,22 +157,23 @@ start-all: .pulled testenv build-wait $(DOCKER_COMPOSE) run -T wait pilosa curl --silent --fail http://pilosa:10101/status $(DOCKER_COMPOSE) run -T wait pilosa-tls curl --silent --cacert /certs/ca.crt --key /certs/theclient.key --cert /certs/theclient.crt --fail https://pilosa-tls:10111/status $(DOCKER_COMPOSE) run -T wait pilosa-auth curl --silent --fail http://pilosa-auth:10105/version + $(DOCKER_COMPOSE) run -T wait dax curl --silent --fail http://dax:8080/computer/status $(DOCKER_COMPOSE) run -T wait kafka nc -z kafka 9092 $(DOCKER_COMPOSE) run -T wait schema-registry curl --silent --fail http://schema-registry:8081/config -start-postgres: build-wait testenv .pulled +start-postgres: build-wait testenv $(DOCKER_COMPOSE) up -d postgres $(DOCKER_COMPOSE) run -T wait postgres pg_isready -h postgres -p 5432 -U postgres -start-pilosa: build-pilosa start-postgres build-wait testenv .pulled +start-pilosa: build-pilosa start-postgres build-wait testenv $(DOCKER_COMPOSE) up -d pilosa $(DOCKER_COMPOSE) run -T wait pilosa curl --silent --fail http://pilosa:10101/status -start-pilosa-tls: build-pilosa-tls build-wait testenv .pulled +start-pilosa-tls: build-pilosa-tls build-wait testenv $(DOCKER_COMPOSE) up -d pilosa-tls $(DOCKER_COMPOSE) run -T wait pilosa-tls curl --silent --cacert /certs/ca.crt --key /certs/theclient.key --cert /certs/theclient.crt --fail https://pilosa-tls:10111/status -start-pilosa-auth: build-pilosa-auth build-wait testenv .pulled +start-pilosa-auth: build-pilosa-auth build-wait testenv $(DOCKER_COMPOSE) up -d pilosa-auth $(DOCKER_COMPOSE) run -T wait pilosa-auth curl --silent --fail http://pilosa-auth:10105/version @@ -191,7 +194,6 @@ startup: start-all shutdown: $(DOCKER_COMPOSE) down -v --remove-orphans - rm -f .pulled test-all: testenv $(MAKE) startup @@ -212,7 +214,7 @@ test-all-kafka-sasl: testenv TCMD ?= ./... # do "make startup", then e.g. "make test-run-local TCMD='-run=MyFavTest ./kafka'" -test-run-local: +test-run-local: vendor pwd $(DOCKER_COMPOSE) build idk-test $(DOCKER_COMPOSE) run -T idk-test go test -mod=vendor -tags=odbc,dynamic $(TCMD) @@ -233,9 +235,6 @@ test-run-kafka-sasl: testenv vendor $(DOCKER_COMPOSE) build idk-test $(DOCKER_COMPOSE) run -T idk-test bash -c "set -o pipefail; go test -v --tags=kafka_sasl -mod=vendor -race -timeout=30m $(TPKG) -covermode=atomic -coverpkg=$(TPKG) -json -coverprofile=/testdata/$(PROJECT)_sasl_coverage.out | tee /testdata/$(PROJECT)_report.out" -.pulled: - $(DOCKER_COMPOSE) pull - touch .pulled testenv: testenv/certs @@ -293,6 +292,9 @@ docker-tag-push: docker push registry.gitlab.com/molecula/featurebase/idk:$(VERSION_ID) @echo Pushed docker image: registry.gitlab.com/molecula/featurebase/idk:$(VERSION_ID) +build-datagen: + CGO_ENABLED=1 $(GO) build -tags dynamic $(GO_BUILD_FLAGS) -o build/datagen ./cmd/datagen + clean: $(MAKE) shutdown rm -rf testenv diff --git a/idk/README.md b/idk/README.md index cb406966a..635bc22ba 100644 --- a/idk/README.md +++ b/idk/README.md @@ -12,19 +12,27 @@ In addition to these dependancies, you will need to be added to the moleculacorp First start the test environment. This is a docker-compose environment that includes pilosa and a confluent kafka stack. Run the following to start those services: - BRANCH_NAME=master make startup + make startup To build and run the integration tests, run: make test-run +An alternative command to use, if you're running tests locally and want human-friendly output, is: + + make test-run-local + +With that command you can also specify individual tests to run like this: + + make test-run-local TCMD='-run=TestJustThisOne .' + Then to shut down the test environment, run: make shutdown You can run all of the previous commands by calling test-all: - BRANCH_NAME=master make test-all + make test-all The previous command is equivalent to running the following: diff --git a/idk/api/schema.go b/idk/api/schema.go index 563429232..d834f9056 100644 --- a/idk/api/schema.go +++ b/idk/api/schema.go @@ -5,8 +5,9 @@ import ( "io" "time" - pilosa "github.com/featurebasedb/featurebase/v3" - pilosaclient "github.com/featurebasedb/featurebase/v3/client" + pilosa "github.com/molecula/featurebase/v3" + pilosaclient "github.com/molecula/featurebase/v3/client" + clienttypes "github.com/molecula/featurebase/v3/client/types" "github.com/pkg/errors" ) @@ -116,7 +117,7 @@ func (f schemaField) applyToPilosa(idx *pilosaclient.Index) error { case f.FieldOptions.EnforceMutualExclusion: opts = []pilosaclient.FieldOption{pilosaclient.OptFieldTypeMutex(pilosaclient.CacheType(f.FieldOptions.CacheType), f.FieldOptions.CacheSize)} case f.FieldOptions.TimeQuantum != "": - opts = []pilosaclient.FieldOption{pilosaclient.OptFieldTypeTime(pilosaclient.TimeQuantum(f.FieldOptions.TimeQuantum))} + opts = []pilosaclient.FieldOption{pilosaclient.OptFieldTypeTime(clienttypes.TimeQuantum(f.FieldOptions.TimeQuantum))} if f.FieldOptions.TTL != "" { ttl, err := time.ParseDuration(f.FieldOptions.TTL) if err != nil { @@ -134,7 +135,7 @@ func (f schemaField) applyToPilosa(idx *pilosaclient.Index) error { case f.FieldOptions.EnforceMutualExclusion: opts = []pilosaclient.FieldOption{pilosaclient.OptFieldTypeMutex(pilosaclient.CacheType(f.FieldOptions.CacheType), f.FieldOptions.CacheSize)} case f.FieldOptions.TimeQuantum != "": - opts = []pilosaclient.FieldOption{pilosaclient.OptFieldTypeTime(pilosaclient.TimeQuantum(f.FieldOptions.TimeQuantum))} + opts = []pilosaclient.FieldOption{pilosaclient.OptFieldTypeTime(clienttypes.TimeQuantum(f.FieldOptions.TimeQuantum))} if f.FieldOptions.TTL != "" { ttl, err := time.ParseDuration(f.FieldOptions.TTL) if err != nil { diff --git a/idk/datagen/cmd.go b/idk/datagen/cmd.go index 53568c818..27fcd3caf 100644 --- a/idk/datagen/cmd.go +++ b/idk/datagen/cmd.go @@ -12,9 +12,10 @@ import ( "sync" "github.com/glycerine/vprint" - pilosaclient "github.com/featurebasedb/featurebase/v3/client" - "github.com/featurebasedb/featurebase/v3/idk" - "github.com/featurebasedb/featurebase/v3/logger" + pilosaclient "github.com/molecula/featurebase/v3/client" + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/idk" + "github.com/molecula/featurebase/v3/logger" "github.com/pkg/errors" "github.com/featurebasedb/featurebase/v3/idk/common" @@ -25,6 +26,7 @@ const ( TargetFeaturebase = "featurebase" TargetKafka = "kafka" TargetKafkaStatic = "kafkastatic" + TargetMDS = "mds" ) // Main is the top-level datagen struct. It represents datagen-specific @@ -41,7 +43,7 @@ type Main struct { cfg SourceGeneratorConfig Source string `short:"s" flag:"source" help:"Source generator type. Running datagen with no arguments will list the available source types."` - Target string `short:"t" flag:"target" help:"Destination for the generated data: [kafka, featurebase]."` + Target string `short:"t" flag:"target" help:"Destination for the generated data: [featurebase, kafka, kafkastatic, mds]."` Concurrency int `short:"c" flag:"concurrency" help:"Number of concurrent sources and indexing routines to launch."` @@ -57,8 +59,10 @@ type Main struct { CustomConfig string `short:"" help:"File from which to pull configuration for 'custom' source."` // Used strictly for configuration of the targets. - Pilosa PilosaConfig - Kafka KafkaConfig + Pilosa PilosaConfig + Kafka KafkaConfig + MDS MDSConfig + FeatureBase FeatureBaseConfig `flag:"featurebase" help:"qualified featurebase table"` DryRun bool `help:"Dry run - just flag parsing."` @@ -77,6 +81,16 @@ type PilosaConfig struct { CacheLength uint64 `help:"Number of batches of ID mappings to cache."` } +// FeatureBaseConfig is meant to represent the scoped (featurebase.*) +// configuration options to be used when target = mds. These are really just a +// sub-set of idk.Main, containing only those arguments that really apply to +// datagen. +type FeatureBaseConfig struct { + OrganizationID string `flag:"org-id" short:"" help:"auto-assigned organization ID"` + DatabaseID string `flag:"db-id" short:"" help:"auto-assigned database ID"` + TableName string `flag:"table-name" short:"" help:"human friendly table name"` +} + // KafkaConfig is meant to represent the scoped (pilosa.*) configuration options // to be used when target = kafka. These are really just a sub-set of kafka.PutSource, // containing only those arguments that really apply to datagen. @@ -89,6 +103,13 @@ type KafkaConfig struct { NumPartitions int `short:"" help:"set partition for kafka cluster"` } +// MDSConfig represents the configuration options to be used when target = mds. +// These are really just a sub-set of idk.Main, containing only those arguments +// that really apply to datagen. +type MDSConfig struct { + Address string `short:"" help:"MDS host:port to connect to"` +} + // NewMain returns a new instance of Main. func NewMain() *Main { return &Main{ @@ -265,6 +286,25 @@ func (m *Main) Preload() error { m.KafkaPut.FBIDField = m.idkMain.IDField m.KafkaPut.FBIndexName = m.Pilosa.Index + case TargetMDS: + m.idkMain.Namespace = "ingester_datagen" + m.idkMain.Concurrency = m.Concurrency + m.idkMain.CacheLength = m.Pilosa.CacheLength + m.idkMain.NewSource = m.newSource + m.idkMain.TrackProgress = m.TrackProgress + m.idkMain.AuthToken = m.AuthToken + m.idkMain.UseShardTransactionalEndpoint = m.UseShardTransactionalEndpoint + if m.Pilosa.BatchSize > 0 { + m.idkMain.BatchSize = m.Pilosa.BatchSize + } + + // MDS-specific + m.idkMain.MDSAddress = m.MDS.Address + m.idkMain.OrganizationID = dax.OrganizationID(m.FeatureBase.OrganizationID) + m.idkMain.DatabaseID = dax.DatabaseID(m.FeatureBase.DatabaseID) + m.idkMain.TableName = dax.TableName(m.FeatureBase.TableName) + m.idkMain.PackBools = "" + default: m.idkMain.Namespace = "ingester_datagen" m.idkMain.Concurrency = m.Concurrency @@ -396,14 +436,15 @@ func (m *Main) PrintPlan() { m.Pilosa.Index = "(not specified)" } fmt.Printf(`Datagen config: - hosts: %s - index: %s - start id: %s - end id: %s - total generated: %s - concurrency: %d - batch size: %s - total batches: %s + hosts: %s + index: %s + start id: %s + end id: %s + total generated: %s + concurrency: %d + batch size: %s + total batches: %s + use shard trans: %v `, strings.Join(Hosts, ", "), m.Pilosa.Index, @@ -413,6 +454,7 @@ func (m *Main) PrintPlan() { concurrency, AddThousandSep(uint64(BatchSize)), AddThousandSep(BatchCount), + m.UseShardTransactionalEndpoint, ) fmt.Println("Schema:") diff --git a/idk/datagen/custom.go b/idk/datagen/custom.go index 444e0fbe0..7c17a4ccf 100644 --- a/idk/datagen/custom.go +++ b/idk/datagen/custom.go @@ -296,6 +296,8 @@ func (c *Custom) Source(cfg SourceConfig) idk.Source { recordsToGenerate = cfg.endAt - cfg.startFrom } return &CustomSource{ + cur: cfg.startFrom, + endAt: cfg.endAt, conf: c.CustomConfig, schema: c.IDKAndGenFields.schema, @@ -333,6 +335,8 @@ var _ idk.Source = (*CustomSource)(nil) // CustomSource is an instance of a source generated // by the Sourcer implementation Custom. type CustomSource struct { + cur uint64 + endAt uint64 conf *CustomConfig generators []FieldGenerator @@ -353,6 +357,10 @@ func (cs *CustomSource) Schema() []idk.Field { // any of the generators are nil, it uses the value from the last // non-nil generator. func (cs *CustomSource) Record() (idk.Record, error) { + if cs.cur >= cs.endAt { + return nil, io.EOF + } + // track last generated value var last interface{} for i, gen := range cs.generators { @@ -364,7 +372,7 @@ func (cs *CustomSource) Record() (idk.Record, error) { var err error cs.record[i], err = gen.Generate(cs.record[i]) if err == io.EOF { - return nil, nil + return nil, io.EOF } if err != nil { return nil, errors.Wrapf(err, "generating for %+v", cs.schema[i]) @@ -377,6 +385,7 @@ func (cs *CustomSource) Record() (idk.Record, error) { } else { cs.recordCounter++ } + cs.cur++ return cs.record, nil } diff --git a/idk/datagen/testdata/basic.yaml b/idk/datagen/testdata/basic.yaml new file mode 100644 index 000000000..1fe8b45c7 --- /dev/null +++ b/idk/datagen/testdata/basic.yaml @@ -0,0 +1,52 @@ +fields: + - name: "an_id" + type: "uint" # (default IDField (non-mutex)) + distribution: "sequential" + min: 0 + max: 268435455 + repeat: false # if false, data generation stops when we hit >= max. only available with sequential + step: 1 + - name: "an_int" + type: "int" # (default IntField) + distribution: "uniform" # uniform or zipfian + min: 0 + max: 500 + - name: "a_random_string" + type: "string" # (default StringField (non-mutex)) + generator_type: "random-string" # used to generate random strings rather than pulling from known set + min_len: 3 + max_len: 3 + charset: "AB" # set of possible characters to pull from when generating random string + - name: "an_id_set" + type: "uint-set" # (default IDArrayField) + min: 0 + max: 1000 + distribution: "uniform" + min_num: 1 + max_num: 6 + - name: "a_string_set" + type: "string-set" # (default StringArrayField) + generator_type: "random-string" # used to generate random strings rather than pulling from known. "distribution" is ignored. + min_len: 4 + max_len: 4 + charset: "0123456789ABCDEF" + min_num: 0 # minimum number of strings in each value (default 0) + max_num: 10 # max number of strings (default to cardinality of source) + +# idk_params describe how data from "fields" should be ingested by IDK +idk_params: + primary_key_config: + field: "id" # if this is a single field named "id" then we'll use uint IDs, if it's empty we'll autogen ids, and if it's anything else we'll do string keys... yes this is a bit hacky, needs to be cleaned up. + # fields is keyed by names of fields from top level "fields". It is + # not required that all fields appear here, those that don't will + # use the default ingestion. + fields: + an_id: + - type: "ID" + mutex: false + name: "id" + a_string_set: + - type: "StringArray" + a_random_string: + - type: "String" + mutex: true diff --git a/idk/datagen/testdata/keyedtable.yaml b/idk/datagen/testdata/keyedtable.yaml new file mode 100644 index 000000000..2b8552477 --- /dev/null +++ b/idk/datagen/testdata/keyedtable.yaml @@ -0,0 +1,58 @@ +fields: + - name: "_id" + type: "string" # (default StringField (non-mutex)) + generator_type: "random-string" # used to generate random strings rather than pulling from known set + min_len: 3 + max_len: 3 + charset: "ABCDEFGH" # set of possible characters to pull from when generating random string + - name: "an_id_set" + type: "uint-set" # (default IDArrayField) + min: 0 + max: 1000 + distribution: "uniform" + min_num: 1 + max_num: 6 + - name: "an_id" + type: "uint" # (default IDArrayField) + min: 0 + max: 1000 + step: 1 + - name: "an_int" + type: "int" # (default IntField) + distribution: "uniform" # uniform or zipfian + min: 20 + max: 45 + - name: "a_string" + type: "string" # (default StringField (non-mutex)) + generator_type: "random-string" # used to generate random strings rather than pulling from known set + min_len: 8 + max_len: 12 + charset: "ZYXWVUTS" # set of possible characters to pull from when generating random string + - name: "a_string_set" + type: "string-set" # (default StringArrayField) + generator_type: "random-string" # used to generate random strings rather than pulling from known. "distribution" is ignored. + min_len: 4 + max_len: 4 + charset: "0123456789ABCDEF" + min_num: 0 # minimum number of strings in each value (default 0) + max_num: 10 # max number of strings (default to cardinality of source) + +# idk_params describe how data from "fields" should be ingested by IDK +idk_params: + primary_key_config: + # if this is a single field named "id" then we'll use uint IDs, if it's + # empty we'll autogen ids, and if it's anything else we'll do string keys... + # yes this is a bit hacky, needs to be cleaned up. + field: "_id" + # fields is keyed by names of fields from top level "fields". It is + # not required that all fields appear here, those that don't will + # use the default ingestion. + fields: + an_id: + - type: "ID" + keyed: false + mutex: true + a_string: + - type: "String" + keyed: true + mutex: true diff --git a/idk/datagen/testdata/keys_ids.yaml b/idk/datagen/testdata/keys_ids.yaml new file mode 100644 index 000000000..5a6ab933f --- /dev/null +++ b/idk/datagen/testdata/keys_ids.yaml @@ -0,0 +1,25 @@ +fields: + - name: "uuid" + type: string + distribution: "shifting" + step: 10 # step for "shifting" distribution means how often do we shift. Every time we generate a value, we'll add an amount to it, every generations, we'll increase the amount by one. + cardinality: 400000000 + max_len: 32 + s: 1.01 + v: 200 + null_chance: 0 + - name: "slice" + type: "uint-set" # (default IDArrayField) + min: 0 + max: 35000 + distribution: "zipfian" + s: 1.1 + v: 5.1 + min_num: 1 + max_num: 50 + +# idk_params describe how data from "fields" should be ingested by IDK +idk_params: + primary_key_config: + field: "uuid" + diff --git a/idk/datagen/testdata/unkeyedtable.yaml b/idk/datagen/testdata/unkeyedtable.yaml new file mode 100644 index 000000000..4bac27cb8 --- /dev/null +++ b/idk/datagen/testdata/unkeyedtable.yaml @@ -0,0 +1,69 @@ +fields: + - name: "rec_id" + type: "uint" # (default IDField (non-mutex)) + distribution: "sequential" + min: 0 + max: 2000000 + repeat: false # if false, data generation stops when we hit >= max. only available with sequential + step: 1 + - name: "an_id" + type: "uint" # (default IDField (non-mutex)) + distribution: "sequential" + min: 0 + max: 2000000 + repeat: false # if false, data generation stops when we hit >= max. only available with sequential + step: 3 + - name: "an_int" + type: "int" # (default IntField) + distribution: "uniform" # uniform or zipfian + min: 20 + max: 45 + - name: "a_string" + type: "string" # (default StringField (non-mutex)) + generator_type: "random-string" # used to generate random strings rather than pulling from known set + min_len: 8 + max_len: 12 + charset: "ZYXWVUTS" # set of possible characters to pull from when generating random string + - name: "an_id_set" + type: "uint-set" # (default IDArrayField) + min: 0 + max: 1000 + distribution: "uniform" + min_num: 1 + max_num: 6 + - name: "a_string_set" + type: "string-set" # (default StringArrayField) + generator_type: "random-string" # used to generate random strings rather than pulling from known. "distribution" is ignored. + min_len: 4 + max_len: 4 + charset: "0123456789ABCDEF" + min_num: 0 # minimum number of strings in each value (default 0) + max_num: 10 # max number of strings (default to cardinality of source) + +# idk_params describe how data from "fields" should be ingested by IDK +idk_params: + primary_key_config: + field: "id" # if this is a single field named "id" then we'll use uint IDs, if it's empty we'll autogen ids, and if it's anything else we'll do string keys... yes this is a bit hacky, needs to be cleaned up. + # fields is keyed by names of fields from top level "fields". It is + # not required that all fields appear here, those that don't will + # use the default ingestion. + fields: + rec_id: + - type: "ID" + mutex: false + keyed: false + name: "id" + an_id: + - type: "ID" + keyed: false + mutex: true + a_string: + - type: "String" + keyed: true + mutex: true + + # - name: "an_id" + #type: "uint" # (default IDArrayField) + #min: 0 + #max: 1000 + #step: 1 diff --git a/idk/docker-compose.yml b/idk/docker-compose.yml index af62a85a7..4b11f955a 100644 --- a/idk/docker-compose.yml +++ b/idk/docker-compose.yml @@ -35,7 +35,7 @@ services: KAFKA_SASL_MECHANISM_INTER_BROKER_PROTOCOL: PLAIN volumes: - ./docker-sasl/ssl_keys:/etc/kafka/secrets - + depends_on: - zookeeper @@ -67,20 +67,22 @@ services: - kafka pilosa: - image: registry.gitlab.com/molecula/featurebase/featurebase:linux-amd64-${BRANCH_NAME} + #image: registry.gitlab.com/molecula/featurebase/featurebase:linux-amd64-${BRANCH_NAME} # (jaffee) I think we should build directly from this tree rather than pulling the image... makes the local dev flow easier + build: .. environment: PILOSA_DATA_DIR: /data PILOSA_BIND: 0.0.0.0:10101 PILOSA_BIND_GRPC: 0.0.0.0:20101 PILOSA_ADVERTISE: pilosa:10101 PILOSA_LOOKUP_DB_DSN: "postgresql://postgres:password@postgres:5432/postgres?sslmode=disable" - depends_on: + depends_on: - postgres volumes: - ./testenv/certs:/certs pilosa-tls: - image: registry.gitlab.com/molecula/featurebase/featurebase:linux-amd64-${BRANCH_NAME} + #image: registry.gitlab.com/molecula/featurebase/featurebase:linux-amd64-${BRANCH_NAME} + build: .. environment: PILOSA_DATA_DIR: /data PILOSA_BIND: https://0.0.0.0:10111 @@ -94,7 +96,8 @@ services: - ./testenv/certs:/certs pilosa-auth: - image: registry.gitlab.com/molecula/featurebase/featurebase:linux-amd64-${BRANCH_NAME} + #image: registry.gitlab.com/molecula/featurebase/featurebase:linux-amd64-${BRANCH_NAME} + build: .. environment: PILOSA_DATA_DIR: /data PILOSA_BIND: 0.0.0.0:10105 @@ -122,8 +125,8 @@ services: - ./docker-sasl/ssl_keys:/ssl_keys - ./testdata:/testdata depends_on: - - kafka - - postgres + #- kafka + #- postgres - fakeidp wait: build: @@ -132,5 +135,21 @@ services: volumes: - ./testenv/certs:/certs - - + dax: + build: + context: .. + dockerfile: Dockerfile-dax + image: idk_dax:latest + environment: + FEATUREBASE_BIND: 0.0.0.0:8080 + FEATUREBASE_VERBOSE: "true" + FEATUREBASE_STORAGE_METHOD: boltdb + FEATUREBASE_STORAGE_DSN: file:/dax-data/mds.boldtb + FEATUREBASE_QUERYER_RUN: "true" + FEATUREBASE_MDS_RUN: "true" + FEATUREBASE_WRITELOGGER_RUN: "true" + FEATUREBASE_WRITELOGGER_CONFIG_DATA_DIR: /dax-data/wl + FEATUREBASE_SNAPSHOTTER_RUN: "true" + FEATUREBASE_SNAPSHOTTER_CONFIG_DATA_DIR: /dax-data/snaps + FEATUREBASE_COMPUTER_RUN: "true" + FEATUREBASE_COMPUTER_CONFIG_DATA_DIR: /dax-data/computer diff --git a/idk/idktest/idktest.go b/idk/idktest/idktest.go index 48473820b..1b4f6fc84 100644 --- a/idk/idktest/idktest.go +++ b/idk/idktest/idktest.go @@ -19,6 +19,7 @@ const ( ) type ExtractResponse struct { + Error string `json:"error"` Results []Result } @@ -72,5 +73,9 @@ func DoExtractQuery(pql, index string) (ExtractResponse, error) { return eResp, errors.Errorf("unmarshaling response: %v", err) } + if eResp.Error != "" { + return eResp, errors.New(eResp.Error) + } + return eResp, nil } diff --git a/idk/ingest.go b/idk/ingest.go index f82232c53..4f04dd352 100644 --- a/idk/ingest.go +++ b/idk/ingest.go @@ -24,15 +24,19 @@ import ( "time" "github.com/felixge/fgprof" - pilosacore "github.com/featurebasedb/featurebase/v3" - pilosagrpc "github.com/featurebasedb/featurebase/v3/api/client" - pilosabatch "github.com/featurebasedb/featurebase/v3/batch" - pilosaclient "github.com/featurebasedb/featurebase/v3/client" - "github.com/featurebasedb/featurebase/v3/logger" - "github.com/featurebasedb/featurebase/v3/pql" - "github.com/featurebasedb/featurebase/v3/prometheus" - proto "github.com/featurebasedb/featurebase/v3/proto" - "github.com/featurebasedb/featurebase/v3/stats" + pilosacore "github.com/molecula/featurebase/v3" + pilosagrpc "github.com/molecula/featurebase/v3/api/client" + pilosabatch "github.com/molecula/featurebase/v3/batch" + pilosaclient "github.com/molecula/featurebase/v3/client" + client_types "github.com/molecula/featurebase/v3/client/types" + "github.com/molecula/featurebase/v3/dax" + mdsclient "github.com/molecula/featurebase/v3/dax/mds/client" + "github.com/molecula/featurebase/v3/idk/mds" + "github.com/molecula/featurebase/v3/logger" + "github.com/molecula/featurebase/v3/pql" + "github.com/molecula/featurebase/v3/prometheus" + proto "github.com/molecula/featurebase/v3/proto" + "github.com/molecula/featurebase/v3/stats" "github.com/pkg/errors" prom "github.com/prometheus/client_golang/prometheus" "github.com/prometheus/client_golang/prometheus/promhttp" @@ -96,6 +100,11 @@ type Main struct { UseShardTransactionalEndpoint bool `flag:"use-shard-transactional-endpoint" help:"Use alternate import endpoint. Currently unstable/testing"` + MDSAddress string `short:"" help:"MDS address."` + OrganizationID dax.OrganizationID `short:"" help:"auto-assigned organization ID"` + DatabaseID dax.DatabaseID `short:"" help:"auto-assigned database ID"` + TableName dax.TableName `short:"" help:"human friendly table name"` + csvWriter *csv.Writer csvFile *os.File // TODO implement the auto-generated IDs... hopefully using Pilosa to manage it. @@ -115,6 +124,11 @@ type Main struct { index *pilosaclient.Index grpcClient *pilosagrpc.GRPCClient + NewImporterFn func() pilosabatch.Importer `flag:"-"` + + SchemaManager SchemaManager `flag:"-"` + Qtbl *dax.QualifiedTable `flag:"-"` + newNexter func(c int) (IDAllocator, error) ra RangeAllocator @@ -212,6 +226,8 @@ func NewMain() *Main { Pprof: "localhost:6062", Stats: "localhost:9093", + SchemaManager: NopSchemaManager, + stats: stats.NopStatsClient, log: logger.NewStandardLogger(os.Stderr), @@ -268,11 +284,15 @@ func (m *Main) run() error { } func (m *Main) clone() (*Main, error) { - schema, err := m.client.Schema() + var index *pilosaclient.Index + + schema, err := m.SchemaManager.Schema() if err != nil { return nil, err } - var index *pilosaclient.Index = schema.Index(m.Index) + + index = schema.Index(m.Index) + // use a copy (schema race condition issues) mClone := *m mClone.index = index @@ -654,7 +674,6 @@ initialFetch: } func (m *Main) Setup() (onFinishRun func(), err error) { - if err := m.validate(); err != nil { return nil, errors.Wrap(err, "validating configuration") } @@ -664,7 +683,7 @@ func (m *Main) Setup() (onFinishRun func(), err error) { fmt.Printf("Delete index '%s' and all data already imported into it (enter '%[1]s' to delete)? ", m.Index) text, _ := reader.ReadString('\n') if strings.TrimSpace(text) == strings.TrimSpace(m.Index) { - err := m.client.DeleteIndex(m.index) + err := m.SchemaManager.DeleteIndex(m.index) if err != nil { return nil, errors.Wrap(err, "deleting index") } @@ -734,11 +753,11 @@ func (m *Main) Setup() (onFinishRun func(), err error) { } if m.AuthToken != "" { - m.AuthToken = "Bearer " + m.AuthToken // Gets added to context - m.client.AuthToken = m.AuthToken // Gets used for calls made from the client + m.AuthToken = "Bearer " + m.AuthToken // Gets added to context + m.SchemaManager.SetAuthToken(m.AuthToken) // Gets used for calls made from the client } - schema, err := m.client.Schema() + schema, err := m.SchemaManager.Schema() if err != nil { return nil, errors.Wrap(err, "getting schema") } @@ -756,10 +775,10 @@ func (m *Main) Setup() (onFinishRun func(), err error) { keyTranslation := len(m.PrimaryKeyFields) > 0 m.index = schema.Index(m.Index, pilosaclient.OptIndexKeys(keyTranslation)) - err = m.client.SyncIndex(m.index) - if err != nil { + if err := m.SchemaManager.SyncIndex(m.index); err != nil { return nil, errors.Wrap(err, "syncing index") } + if m.AutoGenerate { shardWidth := m.index.ShardWidth() if shardWidth == 0 { @@ -953,21 +972,57 @@ func (m *Main) NewLookupClient() (*PostgresClient, error) { func (m *Main) setupClient() (*tls.Config, error) { var tlsConfig *tls.Config var err error + var opts = []pilosaclient.ClientOption{pilosaclient.OptClientStatsClient(m.stats)} if m.TLS.CertificatePath != "" { tlsConfig, err = GetTLSConfig(&m.TLS, m.Log()) if err != nil { return nil, errors.Wrap(err, "getting TLS config") } - m.client, err = pilosaclient.NewClient(m.PilosaHosts, pilosaclient.OptClientTLSConfig(tlsConfig), pilosaclient.OptClientStatsClient(m.stats)) + opts = append(opts, pilosaclient.OptClientTLSConfig(tlsConfig)) + } else { + opts = append(opts, + pilosaclient.OptClientRetries(2), + pilosaclient.OptClientTotalPoolSize(1000), + pilosaclient.OptClientPoolSizePerRoute(400), + ) + } + if m.useMDS() { + opts = append(opts, pilosaclient.OptClientPathPrefix(dax.ServicePrefixComputer)) + } + + m.client, err = pilosaclient.NewClient(m.PilosaHosts, opts...) + if err != nil { + return nil, errors.Wrap(err, "getting featurebase client") + } + + // Set up the SchemaManager + if m.useMDS() { + ctx := context.Background() + + // MDS doesn't auto-create a table based on IDK ingest; the table must + // already exist. + mdsClient := mdsclient.New(dax.Address(m.MDSAddress)) + qual := dax.NewTableQualifier(m.OrganizationID, m.DatabaseID) + qtid, err := mdsClient.TableID(ctx, qual, m.TableName) if err != nil { - return nil, errors.Wrap(err, "getting featurebase client with TLS") + return nil, errors.Wrapf(err, "getting table id: qual: %s, table name: %s", qual, m.TableName) + } + qtbl, err := mdsClient.Table(ctx, qtid) + if err != nil { + return nil, errors.Wrapf(err, "getting table id: qtid: %s", qtid) + } + + m.Qtbl = qtbl + m.Index = string(qtbl.Key()) + m.SchemaManager = mds.NewSchemaManager(dax.Address(m.MDSAddress), qual) + + m.NewImporterFn = func() pilosabatch.Importer { + return mds.NewImporter(mdsClient, qtbl) } } else { - m.client, err = pilosaclient.NewClient(m.PilosaHosts, pilosaclient.OptClientRetries(2), pilosaclient.OptClientTotalPoolSize(1000), pilosaclient.OptClientPoolSizePerRoute(400), pilosaclient.OptClientStatsClient(m.stats)) - if err != nil { - return nil, errors.Wrap(err, "getting featurebase client") - } + m.SchemaManager = m.client } + return tlsConfig, nil } @@ -1127,7 +1182,7 @@ func (m *Main) runDeleter(c int, limitCounter *msgCounter) error { return errors.Wrap(err, "retrieving values for delete") } - trns, err := m.client.StartTransaction("", time.Minute, false, time.Hour) + trns, err := m.SchemaManager.StartTransaction("", time.Minute, false, time.Hour) if err != nil { return errors.Wrap(err, "starting transaction") } @@ -1162,7 +1217,7 @@ func (m *Main) runDeleter(c int, limitCounter *msgCounter) error { // get field, refreshing schema if needed field, ok := index.Fields()[fieldName] if !ok { - schema, err := m.client.Schema() + schema, err := m.SchemaManager.Schema() if err != nil { return errors.Wrap(err, "unknown field, getting new schema") } @@ -1241,7 +1296,7 @@ func (m *Main) runDeleter(c int, limitCounter *msgCounter) error { return errors.Errorf("unhandled field type %s", field.Options().Type()) } } - _, err = m.client.FinishTransaction(trns.ID) + _, err = m.SchemaManager.FinishTransaction(trns.ID) if err != nil { return errors.Wrap(err, "finishing transaction") } @@ -1307,7 +1362,7 @@ func inspect(grpcClient *pilosagrpc.GRPCClient, index string, columnIDs []uint64 } func (m *Main) findPrimary() (*url.URL, error) { - status, err := m.client.Status() + status, err := m.SchemaManager.Status() if err != nil { return nil, errors.Wrap(err, "looking up cluster status") } @@ -1578,7 +1633,7 @@ func (m *Main) batchFromSchema(schema []Field) ([]Recordizer, pilosabatch.Record hasMutex = true opts = append(opts, CacheConfigOf(fld).mutexOption()) } else if q := QuantumOf(fld); q != "" { - opts = append(opts, pilosaclient.OptFieldTypeTime(pilosaclient.TimeQuantum(q))) + opts = append(opts, pilosaclient.OptFieldTypeTime(client_types.TimeQuantum(q))) ttl, err := TTLOf(fld) if err != nil { @@ -1749,20 +1804,29 @@ func (m *Main) batchFromSchema(schema []Field) ([]Recordizer, pilosabatch.Record } } - err = m.client.SyncIndex(m.index) + err = m.SchemaManager.SyncIndex(m.index) if err != nil { - return nil, nil, nil, nil, errors.Wrap(err, "syncing schema") + return nil, nil, nil, nil, errors.Wrap(err, "syncing index") } // Now we need to get the schema back from the server and rewrite // our local fields to make sure we have all the computed options // (like 'base' on int fields). - sSchema, err := m.client.Schema() + sSchema, err := m.SchemaManager.Schema() if err != nil { return nil, nil, nil, nil, errors.Wrap(err, "fetching final schema") } + + // The table name we use to look up in the index in the schema returned from + // SchemaManager.Schema() differs depending on whether this is MDS supported + // or not (i.e. whether it has a table qualifier). + tblName := m.index.Name() + if m.useMDS() { + tblName = string(m.Qtbl.Key()) + } + for i, fld := range fields { - fields[i] = sSchema.Index(m.index.Name()).Field(fld.Name()) + fields[i] = sSchema.Index(tblName).Field(fld.Name()) } batch, err := m.newBatch(fields) @@ -1777,14 +1841,23 @@ func (m *Main) batchFromSchema(schema []Field) ([]Recordizer, pilosabatch.Record } func (m *Main) newBatch(fields []*pilosaclient.Field) (pilosabatch.RecordBatch, error) { - return pilosabatch.NewBatch(pilosaclient.NewImporter(m.client), m.BatchSize, pilosaclient.FromClientIndex(m.index), pilosaclient.FromClientFields(fields), + opts := []pilosabatch.BatchOption{ pilosabatch.OptLogger(m.log), pilosabatch.OptCacheMaxAge(m.CacheLength), pilosabatch.OptSplitBatchMode(m.ExpSplitBatchMode), pilosabatch.OptMaxStaleness(m.BatchMaxStaleness), pilosabatch.OptKeyTranslateBatchSize(m.KeyTranslateBatchSize), pilosabatch.OptUseShardTransactionalEndpoint(m.UseShardTransactionalEndpoint), - ) + } + var importer pilosabatch.Importer + if m.useMDS() { + importer = m.NewImporterFn() + } else { + importer = pilosaclient.NewImporter(m.client) + } + opts = append(opts, pilosabatch.OptImporter(importer)) + + return pilosabatch.NewBatch(importer, m.BatchSize, pilosaclient.FromClientIndex(m.index), pilosaclient.FromClientFields(fields), opts...) } // validateField ensures that the field is configured correctly. @@ -1821,7 +1894,7 @@ func (m *Main) checkFieldCompatibility(pFld *pilosaclient.Field, iFld Field, pac // against pFldOpts. iFldOpts := struct { fieldType pilosaclient.FieldType - timeQuantum pilosaclient.TimeQuantum + timeQuantum client_types.TimeQuantum cacheType pilosaclient.CacheType cacheSize int min pql.Decimal @@ -1860,7 +1933,7 @@ func (m *Main) checkFieldCompatibility(pFld *pilosaclient.Field, iFld Field, pac iFldOpts.cacheSize = pilosacore.DefaultCacheSize } else if q := QuantumOf(fld); q != "" { iFldOpts.fieldType = pilosaclient.FieldTypeTime - iFldOpts.timeQuantum = pilosaclient.TimeQuantum(q) + iFldOpts.timeQuantum = client_types.TimeQuantum(q) ttl, err := TTLOf(fld) if err != nil { @@ -2083,9 +2156,18 @@ func (m *Main) validate() error { return errors.New("must set exactly one of --primary-key-field , --id-field , --auto-generate") } - if m.Index == "" { + if m.useMDS() { + if m.OrganizationID == "" { + return errors.New("must set an organization with --featurebase.org-id") + } else if m.DatabaseID == "" { + return errors.New("must set a database with --featurebase.db-id") + } else if m.TableName == "" { + return errors.New("must set a table with --featurebase.table-name") + } + } else if m.Index == "" { return errors.New("must set an index with --pilosa.index") } + if m.NewSource == nil { return errors.New("must set a NewSource function on IDK ingester") } @@ -2105,3 +2187,7 @@ func (m *Main) allowError(err error) bool { return false } } + +func (m *Main) useMDS() bool { + return m.MDSAddress != "" +} diff --git a/idk/ingest_test.go b/idk/ingest_test.go index 92108c6d4..1ff4cd4f7 100644 --- a/idk/ingest_test.go +++ b/idk/ingest_test.go @@ -19,6 +19,15 @@ import ( "github.com/featurebasedb/featurebase/v3/idk/idktest" "github.com/featurebasedb/featurebase/v3/logger" "github.com/golang-jwt/jwt" + "github.com/molecula/featurebase/v3/authn" + batch "github.com/molecula/featurebase/v3/batch" + pilosaclient "github.com/molecula/featurebase/v3/client" + "github.com/molecula/featurebase/v3/dax" + mdsclient "github.com/molecula/featurebase/v3/dax/mds/client" + "github.com/molecula/featurebase/v3/idk/idktest" + "github.com/molecula/featurebase/v3/idk/mds" + "github.com/molecula/featurebase/v3/logger" + "github.com/molecula/featurebase/v3/pql" "github.com/pkg/errors" "github.com/stretchr/testify/assert" ) @@ -38,6 +47,24 @@ func configureTestFlags(main *Main) { main.Stats = "" } +func configureTestFlagsMDS(main *Main, address dax.Address, qtbl *dax.QualifiedTable) { + main.MDSAddress = address.String() + main.Stats = "" + main.Pprof = "" + main.PackBools = "" + main.OrganizationID = qtbl.Qualifier().OrganizationID + main.DatabaseID = qtbl.Qualifier().DatabaseID + main.TableName = qtbl.Name + main.Qtbl = qtbl + main.SchemaManager = mds.NewSchemaManager(address, qtbl.Qualifier()) + main.Index = string(qtbl.Key()) + + mdsClient := mdsclient.New(dax.Address(address)) + main.NewImporterFn = func() batch.Importer { + return mds.NewImporter(mdsClient, qtbl) + } +} + func TestErrFlush(t *testing.T) { ts := newTestSource([]Field{IDField{NameVal: "aval"}}, [][]interface{}{ @@ -1719,3 +1746,179 @@ func TestBoolIngest(t *testing.T) { }) } } + +func TestBatchTargetMDS(t *testing.T) { + var mdsHost string + if mds, ok := os.LookupEnv("IDK_TEST_MDS_HOST"); ok { + mdsHost = mds + } else { + mdsHost = "dax:8080" + } + + mdsAddress := dax.Address(mdsHost) + orgID := dax.OrganizationID("acme") + dbID := dax.DatabaseID("db1") + + t.Run("FieldTypes", func(t *testing.T) { + tests := []struct { + fieldType dax.FieldType + fieldOptions dax.FieldOptions + fieldFn fieldFn + in [][]interface{} + }{ + // { + // fieldType: types.FieldTypeBool, + // fieldFn: boolFn, + // in: [][]interface{}{ + // {1, true}, + // {2, false}, + // }, + // }, + { + fieldType: dax.FieldTypeDecimal, + fieldOptions: dax.FieldOptions{ + Scale: 4, + }, + fieldFn: decimalFn, + in: [][]interface{}{ + {1, "12.3456"}, + {2, "-7.8"}, + }, + }, + { + fieldType: dax.FieldTypeID, + fieldFn: idFn, + in: [][]interface{}{ + {1, uint64(11)}, + {2, uint64(22)}, + }, + }, + { + fieldType: dax.FieldTypeIDSet, + fieldFn: idSetFn, + in: [][]interface{}{ + {1, []uint64{11, 12, 13}}, + {2, []uint64{22, 24, 26, 28}}, + }, + }, + { + fieldType: dax.FieldTypeInt, + fieldFn: intFn, + in: [][]interface{}{ + {1, int(11)}, + {2, int(-22)}, + }, + fieldOptions: dax.FieldOptions{Min: pql.NewDecimal(-100, 0), Max: pql.NewDecimal(100, 0)}, + }, + { + fieldType: dax.FieldTypeString, + fieldFn: stringFn, + in: [][]interface{}{ + {1, "cycling"}, + {2, "running"}, + }, + }, + { + fieldType: dax.FieldTypeStringSet, + fieldFn: stringSetFn, + in: [][]interface{}{ + {1, []string{"cycling", "swimming"}}, + {2, []string{"running", "cooking"}}, + }, + }, + { + fieldType: dax.FieldTypeTimestamp, + fieldFn: timestampFn, + in: [][]interface{}{ + {1, time.Now()}, + {2, time.Now()}, + }, + fieldOptions: dax.FieldOptions{TimeUnit: "s", Epoch: time.Unix(0, 0)}, // TODO w/o this, it used to silently fail (just logs in the mds svc). Eventually should probably have a default unit and not fail at all. + }, + } + for i, test := range tests { + t.Run(fmt.Sprintf("test-%s-%d", test.fieldType, i), func(t *testing.T) { + // Generate a random tableName to use for the test. + rand.Seed(time.Now().UTC().UnixNano()) + tableName := fmt.Sprintf("tbl_%s_%d", test.fieldType, rand.Intn(100000)) + fieldName := fmt.Sprintf("fld_%s_%d", test.fieldType, rand.Intn(100000)) + + tblName := dax.TableName(tableName) + + tbl := dax.NewTable(tblName) + tbl.PartitionN = dax.DefaultPartitionN + tbl.Fields = []*dax.Field{ + { + Name: dax.PrimaryKeyFieldName, + Type: dax.FieldTypeID, + }, + { + Name: dax.FieldName(fieldName), + Type: test.fieldType, + Options: test.fieldOptions, + }, + } + + qtbl := dax.NewQualifiedTable( + dax.NewTableQualifier(orgID, dbID), + tbl, + ) + + ctx := context.Background() + + // Create the table in MDS Schemar. + mdsClient := mdsclient.New(mdsAddress) + if err := mdsClient.CreateTable(ctx, qtbl); err != nil { + t.Fatalf("creating table: %v", err) + } + + qtblWithID, err := mdsClient.Table(ctx, qtbl.QualifiedID()) + assert.NoError(t, err) + + ts := newTestSource([]Field{ + IDField{NameVal: "id"}, + test.fieldFn(fieldName, test.fieldOptions), + }, test.in) + + ingester := NewMain() + configureTestFlagsMDS(ingester, mdsAddress, qtblWithID) + + ingester.NewSource = func() (Source, error) { return ts, nil } + ingester.BatchSize = 10 + //ingester.PrimaryKeyFields = []string{"rcid"} + ingester.IDField = "id" + + if err := ingester.Run(); err != nil { + t.Fatalf("running ingester: %v", err) + } + }) + } + }) +} + +type fieldFn func(string, dax.FieldOptions) Field + +func boolFn(name string, fo dax.FieldOptions) Field { + return BoolField{NameVal: name} +} +func decimalFn(name string, fo dax.FieldOptions) Field { + return DecimalField{NameVal: name, Scale: fo.Scale} +} +func idFn(name string, fo dax.FieldOptions) Field { + return IDField{NameVal: name, Mutex: true} +} +func idSetFn(name string, fo dax.FieldOptions) Field { + return IDArrayField{NameVal: name} +} +func intFn(name string, fo dax.FieldOptions) Field { + return IntField{NameVal: name} +} +func stringFn(name string, fo dax.FieldOptions) Field { + return StringField{NameVal: name, Mutex: true} +} +func stringSetFn(name string, fo dax.FieldOptions) Field { + return StringArrayField{NameVal: name} +} +func timestampFn(name string, fo dax.FieldOptions) Field { + return TimestampField{NameVal: name} +} diff --git a/idk/interfaces.go b/idk/interfaces.go index 950a651e0..6e24ce29c 100644 --- a/idk/interfaces.go +++ b/idk/interfaces.go @@ -1309,3 +1309,51 @@ func (f Fields) ContainsBool() bool { } return false } + +// SchemaManager is meant to be an interface for managing schema information; +// i.e. for interacting with a single source of truth for schema information, +// like the MDS Schemar. But... it currently contains methods which are not +// related to schema because the first goal was just to introduce an interface +// in ingest.go for any methods being called on *m.client. We don't want a +// FeatureBase client directly called from ingest, rather, we want to call these +// interface methods and allow for different implementations (such as an MDS +// implementation which uses the Schemar in MDS as opposed to a FeatureBase node +// or cluster). +type SchemaManager interface { + StartTransaction(id string, timeout time.Duration, exclusive bool, requestTimeout time.Duration) (*pilosacore.Transaction, error) + FinishTransaction(id string) (*pilosacore.Transaction, error) + Schema() (*pilosaclient.Schema, error) + SyncIndex(index *pilosaclient.Index) error + DeleteIndex(index *pilosaclient.Index) error + Status() (pilosaclient.Status, error) + SetAuthToken(string) +} + +// Ensure type implements interface. +var _ SchemaManager = &nopSchemaManager{} + +// NopSchemaManager is an implementation of the SchemaManager interface that +// doesn't do anything. +var NopSchemaManager SchemaManager = &nopSchemaManager{} + +type nopSchemaManager struct{} + +func (n *nopSchemaManager) StartTransaction(id string, timeout time.Duration, exclusive bool, requestTimeout time.Duration) (*pilosacore.Transaction, error) { + return nil, nil +} +func (n *nopSchemaManager) FinishTransaction(id string) (*pilosacore.Transaction, error) { + return nil, nil +} +func (n *nopSchemaManager) Schema() (*pilosaclient.Schema, error) { + return nil, nil +} +func (n *nopSchemaManager) SyncIndex(index *pilosaclient.Index) error { + return nil +} +func (n *nopSchemaManager) DeleteIndex(index *pilosaclient.Index) error { + return nil +} +func (n *nopSchemaManager) Status() (pilosaclient.Status, error) { + return pilosaclient.Status{}, nil +} +func (n *nopSchemaManager) SetAuthToken(token string) {} diff --git a/idk/mds/importer.go b/idk/mds/importer.go new file mode 100644 index 000000000..06737274a --- /dev/null +++ b/idk/mds/importer.go @@ -0,0 +1,276 @@ +package mds + +import ( + "context" + "sync" + "time" + + featurebase "github.com/molecula/featurebase/v3" + "github.com/molecula/featurebase/v3/batch" + featurebaseclient "github.com/molecula/featurebase/v3/client" + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/dax/mds/controller/partitioner" + "github.com/molecula/featurebase/v3/roaring" + "github.com/pkg/errors" +) + +// Ensure type implements interface. +var _ batch.Importer = &importer{} + +// importer +type importer struct { + mds MDS + + mu sync.Mutex + qtbl *dax.QualifiedTable +} + +func NewImporter(mds MDS, qtbl *dax.QualifiedTable) *importer { + return &importer{ + mds: mds, + qtbl: qtbl, + } +} + +// fbClient currently returns a new FeatureBase client (for the address) for +// every call to this method. We could cache these connections in a map (keyed +// on address) to avoid creating a new client for an address that we already +// have a client for. +func (m *importer) fbClient(address dax.Address) (*featurebaseclient.Client, error) { + // Set up a FeatureBase client with address. + return featurebaseclient.NewClient(address.String(), + featurebaseclient.OptClientRetries(2), + featurebaseclient.OptClientTotalPoolSize(1000), + featurebaseclient.OptClientPoolSizePerRoute(400), + featurebaseclient.OptClientPathPrefix(dax.ServicePrefixComputer), + //featurebaseclient.OptClientStatsClient(m.stats), + ) +} + +func (m *importer) StartTransaction(ctx context.Context, id string, timeout time.Duration, exclusive bool, requestTimeout time.Duration) (*featurebase.Transaction, error) { + return &featurebase.Transaction{ + ID: "not-used", + }, nil +} + +func (m *importer) FinishTransaction(ctx context.Context, id string) (*featurebase.Transaction, error) { + return nil, nil +} + +func (m *importer) CreateIndexKeys(ctx context.Context, idx *featurebase.IndexInfo, keys ...string) (map[string]uint64, error) { + qtbl, err := m.getQtbl(ctx, idx.Name) + if err != nil { + return nil, errors.Wrapf(err, "getting qtbl") + } + + out := make(map[string]uint64) + + partitioner := partitioner.NewPartitioner() + + // Get the partition for each key (map[int][]string). + partitions := partitioner.PartitionsForKeys(qtbl.Key(), qtbl.PartitionN, keys...) + + // TODO: we can be more efficient here by calling IngestPartitions() with + // all the partitions at once, then getting the distinct list of addresses + // and looping over that instead. + for partition, ks := range partitions { + address, err := m.mds.IngestPartition(context.Background(), qtbl.QualifiedID(), partition) + if err != nil { + return nil, errors.Wrapf(err, "calling ingest-partition on table: %s, partition: %d", qtbl, partition) + } + + fbClient, err := m.fbClient(address) + if err != nil { + return nil, errors.Wrap(err, "getting featurebase client") + } + + stringToIDMap, err := fbClient.CreateIndexKeys(featurebaseclient.ToClientIndex(idx), ks...) + if err != nil { + return nil, errors.Wrapf(err, "creating index keys for partition: %d", partition) + } + + for str, id := range stringToIDMap { + out[str] = id + } + } + + return out, nil +} + +func (m *importer) CreateFieldKeys(ctx context.Context, index string, field *featurebase.FieldInfo, keys ...string) (map[string]uint64, error) { + qtbl, err := m.getQtbl(ctx, index) + if err != nil { + return nil, errors.Wrapf(err, "getting qtbl") + } + + // For now, we are going to direct all field key translation to the same + // node handling index key translation for the table, partition 0. + // TODO: we should be able to partition field key translation on fieldName. + // If we do that, we might also consider whether we want to support a + // different partitionN for field translation. + partition := dax.PartitionNum(0) + + address, err := m.mds.IngestPartition(context.Background(), qtbl.QualifiedID(), partition) + if err != nil { + return nil, errors.Wrapf(err, "calling ingest-partition on table: %s, partition: %d", qtbl, partition) + } + + // Set up a FeatureBase client with address. + fbClient, err := m.fbClient(address) + if err != nil { + return nil, errors.Wrap(err, "getting featurebase client") + } + + cfld, err := featurebaseclient.ToClientField(index, field) + if err != nil { + return nil, errors.Wrap(err, "converting fieldinfo to client field") + } + + return fbClient.CreateFieldKeys(cfld, keys...) +} + +func (m *importer) ImportRoaringBitmap(ctx context.Context, index string, field *featurebase.FieldInfo, shard uint64, views map[string]*roaring.Bitmap, clear bool) error { + qtbl, err := m.getQtbl(ctx, index) + if err != nil { + return errors.Wrapf(err, "getting qtbl") + } + + address, err := m.mds.IngestShard(context.Background(), qtbl.QualifiedID(), dax.ShardNum(shard)) + if err != nil { + return errors.Wrap(err, "calling ingest-shard") + } + + // Set up a FeatureBase client with address. + fbClient, err := m.fbClient(address) + if err != nil { + return errors.Wrap(err, "getting featurebase client") + } + + cfld, err := featurebaseclient.ToClientField(index, field) + if err != nil { + return errors.Wrap(err, "converting fieldinfo to client field") + } + + return fbClient.ImportRoaringBitmap(cfld, shard, views, clear) +} + +func (m *importer) ImportRoaringShard(ctx context.Context, index string, shard uint64, request *featurebase.ImportRoaringShardRequest) error { + qtbl, err := m.getQtbl(ctx, index) + if err != nil { + return errors.Wrapf(err, "getting qtbl") + } + + address, err := m.mds.IngestShard(context.Background(), qtbl.QualifiedID(), dax.ShardNum(shard)) + if err != nil { + return errors.Wrap(err, "calling ingest-shard") + } + + // Set up a FeatureBase client with address. + fbClient, err := m.fbClient(address) + if err != nil { + return errors.Wrap(err, "getting featurebase client") + } + + return fbClient.ImportRoaringShard(index, shard, request) +} + +func (m *importer) EncodeImportValues(ctx context.Context, index string, field *featurebase.FieldInfo, shard uint64, vals []int64, ids []uint64, clear bool) (path string, data []byte, err error) { + qtbl, err := m.getQtbl(ctx, index) + if err != nil { + return "", nil, errors.Wrapf(err, "getting qtbl") + } + + address, err := m.mds.IngestShard(context.Background(), qtbl.QualifiedID(), dax.ShardNum(shard)) + if err != nil { + return "", nil, errors.Wrap(err, "calling ingest-shard") + } + + // Set up a FeatureBase client with address. + fbClient, err := m.fbClient(address) + if err != nil { + return "", nil, errors.Wrap(err, "getting featurebase client") + } + + cfld, err := featurebaseclient.ToClientField(index, field) + if err != nil { + return "", nil, errors.Wrap(err, "converting fieldinfo to client field") + } + + return fbClient.EncodeImportValues(cfld, shard, vals, ids, clear) +} + +func (m *importer) EncodeImport(ctx context.Context, index string, field *featurebase.FieldInfo, shard uint64, vals, ids []uint64, clear bool) (path string, data []byte, err error) { + qtbl, err := m.getQtbl(ctx, index) + if err != nil { + return "", nil, errors.Wrapf(err, "getting qtbl") + } + + address, err := m.mds.IngestShard(context.Background(), qtbl.QualifiedID(), dax.ShardNum(shard)) + if err != nil { + return "", nil, errors.Wrap(err, "calling ingest-shard") + } + + // Set up a FeatureBase client with address. + fbClient, err := m.fbClient(address) + if err != nil { + return "", nil, errors.Wrap(err, "getting featurebase client") + } + + cfld, err := featurebaseclient.ToClientField(index, field) + if err != nil { + return "", nil, errors.Wrap(err, "converting fieldinfo to client field") + } + + return fbClient.EncodeImport(cfld, shard, vals, ids, clear) +} + +func (m *importer) DoImport(ctx context.Context, index string, field *featurebase.FieldInfo, shard uint64, path string, data []byte) error { + qtbl, err := m.getQtbl(ctx, index) + if err != nil { + return errors.Wrapf(err, "getting qtbl") + } + + address, err := m.mds.IngestShard(context.Background(), qtbl.QualifiedID(), dax.ShardNum(shard)) + if err != nil { + return errors.Wrap(err, "calling ingest-shard") + } + + // Set up a FeatureBase client with address. + fbClient, err := m.fbClient(address) + if err != nil { + return errors.Wrap(err, "getting featurebase client") + } + + return fbClient.DoImport(index, shard, path, data) +} + +func (m *importer) StatsTiming(name string, value time.Duration, rate float64) {} + +// getQtbl takes a table (TableKey) and sets the local m.qtbl value. When we +// originally set up this type, it was only used by IDK, and the table was known +// at the beginning of the process, so it could be set on this import. But +// later, we used this importer in the Queryer, and it gets set up prior to +// parsing the table out of sql; which means we don't know what the table is +// yet. So this method allows us to use the table which is passed into each +// method to determine the table. We look it up from mds schema once and save it +// in m.qtbl for any further method calls. +func (m *importer) getQtbl(ctx context.Context, table string) (*dax.QualifiedTable, error) { + m.mu.Lock() + defer m.mu.Unlock() + + if m.qtbl != nil { + return m.qtbl, nil + } + + tkey := dax.TableKey(table) + qtid := tkey.QualifiedTableID() + + qtbl, err := m.mds.Table(ctx, qtid) + if err != nil { + return nil, errors.Wrap(err, "getting table") + } + + m.qtbl = qtbl + + return qtbl, nil +} diff --git a/idk/mds/mds.go b/idk/mds/mds.go new file mode 100644 index 000000000..e81f69914 --- /dev/null +++ b/idk/mds/mds.go @@ -0,0 +1,21 @@ +// Package mds contains the implementation of the SchemaManager interface. +package mds + +import ( + "context" + + "github.com/molecula/featurebase/v3/dax" +) + +// MDS represents the MDS methods which importer uses. +type MDS interface { + IngestPartition(ctx context.Context, qtid dax.QualifiedTableID, partition dax.PartitionNum) (dax.Address, error) + IngestShard(ctx context.Context, qtid dax.QualifiedTableID, shard dax.ShardNum) (dax.Address, error) + + // Table was added so the `importer` instance (in this package) of the + // batch.Importer interface could lookup up a table based on the name + // provided in a method, as opposed to setting the table up front. This is + // because in queryer, we don't know the table yet, because we haven't + // parsed the sql yet. + Table(ctx context.Context, qtid dax.QualifiedTableID) (*dax.QualifiedTable, error) +} diff --git a/idk/mds/schemamanager.go b/idk/mds/schemamanager.go new file mode 100644 index 000000000..22ba54ef9 --- /dev/null +++ b/idk/mds/schemamanager.go @@ -0,0 +1,114 @@ +package mds + +import ( + "context" + "time" + + featurebase "github.com/molecula/featurebase/v3" + featurebase_client "github.com/molecula/featurebase/v3/client" + "github.com/molecula/featurebase/v3/dax" + mdsclient "github.com/molecula/featurebase/v3/dax/mds/client" + "github.com/molecula/featurebase/v3/errors" +) + +// Ensure type implements interface. +// var _ idk.SchemaManager = &schemaManager{} + +// schemaManager +type schemaManager struct { + client *mdsclient.Client + qual dax.TableQualifier +} + +func NewSchemaManager(mdsAddress dax.Address, qual dax.TableQualifier) *schemaManager { + return &schemaManager{ + client: mdsclient.New(mdsAddress), + qual: qual, + } +} + +func (s *schemaManager) StartTransaction(id string, timeout time.Duration, exclusive bool, requestTimeout time.Duration) (*featurebase.Transaction, error) { + return nil, nil +} +func (s *schemaManager) FinishTransaction(id string) (*featurebase.Transaction, error) { + return nil, nil +} +func (s *schemaManager) Schema() (*featurebase_client.Schema, error) { + // Create a temp schema object to mimic what the FeatureBase client Schema() + // method returns. + schema := featurebase_client.NewSchema() + + tables, err := s.client.Tables(context.Background(), s.qual) + if err != nil { + return nil, err + } + + for _, qtbl := range tables { + idx := schema.Index(string(qtbl.Key()), featurebase_client.OptIndexKeys(qtbl.StringKeys())) + for _, fld := range qtbl.Fields { + opts := make([]featurebase_client.FieldOption, 0) + + switch fld.Type { + case dax.FieldTypeBool: + opts = append(opts, featurebase_client.OptFieldTypeBool()) + case dax.FieldTypeDecimal: + opts = append(opts, featurebase_client.OptFieldTypeDecimal( + fld.Options.Scale, + )) + case dax.FieldTypeID: + opts = append(opts, featurebase_client.OptFieldTypeMutex( + featurebase_client.CacheType(fld.Options.CacheType), + int(fld.Options.CacheSize), + )) + case dax.FieldTypeIDSet: + opts = append(opts, featurebase_client.OptFieldTypeSet( + featurebase_client.CacheType(fld.Options.CacheType), + int(fld.Options.CacheSize), + )) + case dax.FieldTypeInt: + opts = append(opts, featurebase_client.OptFieldTypeInt( + fld.Options.Min.ToInt64(0), + fld.Options.Max.ToInt64(0), + )) + case dax.FieldTypeString: + opts = append(opts, + featurebase_client.OptFieldTypeMutex( + featurebase_client.CacheType(fld.Options.CacheType), + int(fld.Options.CacheSize), + ), + featurebase_client.OptFieldKeys(true), + ) + case dax.FieldTypeStringSet: + opts = append(opts, + featurebase_client.OptFieldTypeSet( + featurebase_client.CacheType(fld.Options.CacheType), + int(fld.Options.CacheSize), + ), + featurebase_client.OptFieldKeys(true), + ) + case dax.FieldTypeTimestamp: + opts = append(opts, featurebase_client.OptFieldTypeTimestamp( + featurebase_client.DefaultEpoch, + fld.Options.TimeUnit, + )) + + default: + return nil, errors.Errorf("unsupported field type: %s (%s)", fld.Name, fld.Type) + } + + _ = idx.Field(string(fld.Name), opts...) + } + } + + return schema, nil +} +func (s *schemaManager) SyncIndex(index *featurebase_client.Index) error { + return nil +} +func (s *schemaManager) DeleteIndex(index *featurebase_client.Index) error { + return nil +} +func (s *schemaManager) Status() (featurebase_client.Status, error) { + return featurebase_client.Status{}, nil +} +func (s *schemaManager) SetAuthToken(token string) {} diff --git a/idk/testdata/basic.yaml b/idk/testdata/basic.yaml new file mode 100644 index 000000000..aca6bfaa5 --- /dev/null +++ b/idk/testdata/basic.yaml @@ -0,0 +1,52 @@ +fields: + - name: "an_id" + type: "uint" # (default IDField (non-mutex)) + distribution: "sequential" + min: 0 + max: 2000000 + repeat: false # if false, data generation stops when we hit >= max. only available with sequential + step: 1 + - name: "an_int" + type: "int" # (default IntField) + distribution: "uniform" # uniform or zipfian + min: 0 + max: 500 + - name: "a_random_string" + type: "string" # (default StringField (non-mutex)) + generator_type: "random-string" # used to generate random strings rather than pulling from known set + min_len: 3 + max_len: 3 + charset: "AB" # set of possible characters to pull from when generating random string + - name: "an_id_set" + type: "uint-set" # (default IDArrayField) + min: 0 + max: 1000 + distribution: "uniform" + min_num: 1 + max_num: 6 + - name: "a_string_set" + type: "string-set" # (default StringArrayField) + generator_type: "random-string" # used to generate random strings rather than pulling from known. "distribution" is ignored. + min_len: 4 + max_len: 4 + charset: "0123456789ABCDEF" + min_num: 0 # minimum number of strings in each value (default 0) + max_num: 10 # max number of strings (default to cardinality of source) + +# idk_params describe how data from "fields" should be ingested by IDK +idk_params: + primary_key_config: + field: "id" # if this is a single field named "id" then we'll use uint IDs, if it's empty we'll autogen ids, and if it's anything else we'll do string keys... yes this is a bit hacky, needs to be cleaned up. + # fields is keyed by names of fields from top level "fields". It is + # not required that all fields appear here, those that don't will + # use the default ingestion. + fields: + an_id: + - type: "ID" + mutex: false + name: "id" + a_string_set: + - type: "StringArray" + a_random_string: + - type: "String" + mutex: true diff --git a/idk/testdata/keyedtable.yaml b/idk/testdata/keyedtable.yaml new file mode 100644 index 000000000..2b8552477 --- /dev/null +++ b/idk/testdata/keyedtable.yaml @@ -0,0 +1,58 @@ +fields: + - name: "_id" + type: "string" # (default StringField (non-mutex)) + generator_type: "random-string" # used to generate random strings rather than pulling from known set + min_len: 3 + max_len: 3 + charset: "ABCDEFGH" # set of possible characters to pull from when generating random string + - name: "an_id_set" + type: "uint-set" # (default IDArrayField) + min: 0 + max: 1000 + distribution: "uniform" + min_num: 1 + max_num: 6 + - name: "an_id" + type: "uint" # (default IDArrayField) + min: 0 + max: 1000 + step: 1 + - name: "an_int" + type: "int" # (default IntField) + distribution: "uniform" # uniform or zipfian + min: 20 + max: 45 + - name: "a_string" + type: "string" # (default StringField (non-mutex)) + generator_type: "random-string" # used to generate random strings rather than pulling from known set + min_len: 8 + max_len: 12 + charset: "ZYXWVUTS" # set of possible characters to pull from when generating random string + - name: "a_string_set" + type: "string-set" # (default StringArrayField) + generator_type: "random-string" # used to generate random strings rather than pulling from known. "distribution" is ignored. + min_len: 4 + max_len: 4 + charset: "0123456789ABCDEF" + min_num: 0 # minimum number of strings in each value (default 0) + max_num: 10 # max number of strings (default to cardinality of source) + +# idk_params describe how data from "fields" should be ingested by IDK +idk_params: + primary_key_config: + # if this is a single field named "id" then we'll use uint IDs, if it's + # empty we'll autogen ids, and if it's anything else we'll do string keys... + # yes this is a bit hacky, needs to be cleaned up. + field: "_id" + # fields is keyed by names of fields from top level "fields". It is + # not required that all fields appear here, those that don't will + # use the default ingestion. + fields: + an_id: + - type: "ID" + keyed: false + mutex: true + a_string: + - type: "String" + keyed: true + mutex: true diff --git a/idk/testdata/unkeyedtable.yaml b/idk/testdata/unkeyedtable.yaml new file mode 100644 index 000000000..4bac27cb8 --- /dev/null +++ b/idk/testdata/unkeyedtable.yaml @@ -0,0 +1,69 @@ +fields: + - name: "rec_id" + type: "uint" # (default IDField (non-mutex)) + distribution: "sequential" + min: 0 + max: 2000000 + repeat: false # if false, data generation stops when we hit >= max. only available with sequential + step: 1 + - name: "an_id" + type: "uint" # (default IDField (non-mutex)) + distribution: "sequential" + min: 0 + max: 2000000 + repeat: false # if false, data generation stops when we hit >= max. only available with sequential + step: 3 + - name: "an_int" + type: "int" # (default IntField) + distribution: "uniform" # uniform or zipfian + min: 20 + max: 45 + - name: "a_string" + type: "string" # (default StringField (non-mutex)) + generator_type: "random-string" # used to generate random strings rather than pulling from known set + min_len: 8 + max_len: 12 + charset: "ZYXWVUTS" # set of possible characters to pull from when generating random string + - name: "an_id_set" + type: "uint-set" # (default IDArrayField) + min: 0 + max: 1000 + distribution: "uniform" + min_num: 1 + max_num: 6 + - name: "a_string_set" + type: "string-set" # (default StringArrayField) + generator_type: "random-string" # used to generate random strings rather than pulling from known. "distribution" is ignored. + min_len: 4 + max_len: 4 + charset: "0123456789ABCDEF" + min_num: 0 # minimum number of strings in each value (default 0) + max_num: 10 # max number of strings (default to cardinality of source) + +# idk_params describe how data from "fields" should be ingested by IDK +idk_params: + primary_key_config: + field: "id" # if this is a single field named "id" then we'll use uint IDs, if it's empty we'll autogen ids, and if it's anything else we'll do string keys... yes this is a bit hacky, needs to be cleaned up. + # fields is keyed by names of fields from top level "fields". It is + # not required that all fields appear here, those that don't will + # use the default ingestion. + fields: + rec_id: + - type: "ID" + mutex: false + keyed: false + name: "id" + an_id: + - type: "ID" + keyed: false + mutex: true + a_string: + - type: "String" + keyed: true + mutex: true + + # - name: "an_id" + #type: "uint" # (default IDArrayField) + #min: 0 + #max: 1000 + #step: 1 diff --git a/index.go b/index.go index be377d6c5..0a6205204 100644 --- a/index.go +++ b/index.go @@ -13,11 +13,12 @@ import ( "sync" "time" - "github.com/featurebasedb/featurebase/v3/disco" - "github.com/featurebasedb/featurebase/v3/pql" - "github.com/featurebasedb/featurebase/v3/roaring" - "github.com/featurebasedb/featurebase/v3/stats" - "github.com/featurebasedb/featurebase/v3/testhook" + "github.com/molecula/featurebase/v3/dax" + "github.com/molecula/featurebase/v3/disco" + "github.com/molecula/featurebase/v3/pql" + "github.com/molecula/featurebase/v3/roaring" + "github.com/molecula/featurebase/v3/stats" + "github.com/molecula/featurebase/v3/testhook" "github.com/pkg/errors" "golang.org/x/sync/errgroup" ) @@ -49,7 +50,8 @@ type Index struct { holder *Holder // Per-partition translation stores - translateStores map[int]TranslateStore + translatePartitions dax.Partitions + translateStores map[int]TranslateStore translationSyncer TranslationSyncer @@ -240,6 +242,28 @@ func (i *Index) open(idx *disco.Index) (err error) { var g errgroup.Group var mu sync.Mutex + // TODO(tlt): this for loop doesn't work because if we assign a + // translate partition to this node later (after the table has been + // created with a sub-set of translatePartitions), then the new + // TranslateStores don't get initialized. For now I just put it back so + // it opens a TranslateStore for every partition no matter what, but we + // really need to have the ApplyDirective logic able to initialize any + // TranslateStore which doesn't already exist (and perhaps shut down any + // that are to be removed). + // + // for _, partition := range i.translatePartitions { + // partitionID := int(partition.Num) + // + // + // TODO(tlt): instead of i.holder.partitionN, we need to use + // len(i.translatePartitions), or actually we need to know the + // keypartitions for the qtbl (i don't think we can rely on the length + // of this slice) but that will only apply here... we need to go through + // all the code and see where these are being used: + // - i.holder.partitionN + // - DefaultPartitionN + // + // for partitionID := 0; partitionID < i.holder.partitionN; partitionID++ { partitionID := partitionID @@ -935,6 +959,21 @@ func (i *Index) DeleteField(name string) error { return i.translationSyncer.Reset() } +// SetTranslatePartitions sets the cached value: translatePartitions. +// +// There's already logic in api_directive.go which creates a new index with +// partitions. This particular function is used when the index already exists on +// the node, but we get a Directive which changes its partition list. In that +// case, we need to update this cached value. Really, this is kind of hacky and +// we need to revisit the ApplyDirective logic so that it's more intuitive with +// respect to index.translatePartitions. +func (i *Index) SetTranslatePartitions(tp dax.Partitions) { + i.mu.Lock() + defer i.mu.Unlock() + + i.translatePartitions = tp +} + type indexSlice []*Index func (p indexSlice) Swap(i, j int) { p[i], p[j] = p[j], p[i] } diff --git a/interfaces.go b/interfaces.go new file mode 100644 index 000000000..c5bb1a104 --- /dev/null +++ b/interfaces.go @@ -0,0 +1,33 @@ +package pilosa + +import ( + "context" + "io" + + "github.com/molecula/featurebase/v3/dax" +) + +// MDS represents the MDS methods which Computer uses. These are typically +// implemented by both the MDS service and the MSD client. +type MDS interface { + RegisterNode(ctx context.Context, node *dax.Node) error + CheckInNode(ctx context.Context, node *dax.Node) error +} + +// WriteLogger represents the WriteLogger methods which Computer uses. These are +// typically implemented by both the WriteLogger service and the WriteLogger +// client. +type WriteLogger interface { + AppendMessage(bucket string, key string, version int, msg []byte) error + LogReader(bucket string, key string, version int) (io.Reader, io.Closer, error) + DeleteLog(bucket string, key string, version int) error +} + +// Snapshotter represents the Snapshotter methods which Computer uses. These are +// typically implemented by both the Snapshotter service and the Snapshotter +// client. +type Snapshotter interface { + Read(bucket string, key string, version int) (io.ReadCloser, error) + Write(bucket string, key string, version int, rc io.ReadCloser) error + WriteTo(bucket string, key string, version int, wrTo io.WriterTo) error +} diff --git a/internal/clustertests/cluster_test.go b/internal/clustertests/cluster_test.go index 56d5db3b8..529b0b0fd 100644 --- a/internal/clustertests/cluster_test.go +++ b/internal/clustertests/cluster_test.go @@ -22,7 +22,6 @@ import ( "github.com/featurebasedb/featurebase/v3/disco" "github.com/featurebasedb/featurebase/v3/logger" "github.com/pkg/errors" - "golang.org/x/sync/errgroup" ) // container turns a docker-compose service name into a container ID @@ -315,91 +314,6 @@ func ingestRandomData(ctx context.Context, cli *pilosa.InternalClient, index, fi return nil } -func TestRetryLogic(t *testing.T) { - if os.Getenv("ENABLE_PILOSA_CLUSTER_TESTS") != "1" { - t.Skip("pilosa cluster tests are not enabled") - } - ctx := context.Background() - auth := false - if os.Getenv("ENABLE_AUTH") == "1" { - auth = true - } - if auth { - token := GetAuthToken(t) - ctx = authn.WithAccessToken(ctx, "Bearer "+token) - } - - var addrs = []string{"pilosa1:10101", "pilosa2:10101", "pilosa3:10101"} - cli, err := getClients(addrs) - if err != nil { - t.Fatalf("getting client: %v", err) - } - g := new(errgroup.Group) - g.Go(func() error { - return ingestRandomData(ctx, cli[0], "testidx1", "testfield1", 100000) - }) - if err := pauseNode(t, "pilosa2"); err != nil { - t.Fatalf("sending pause command: %v", err) - } - if err := pauseNode(t, "pilosa3"); err != nil { - t.Fatalf("sending pause command: %v", err) - } - time.Sleep(6 * time.Second) - if err := unpauseNode(t, "pilosa2"); err != nil { - t.Fatalf("sending pause command: %v", err) - } - if err := unpauseNode(t, "pilosa3"); err != nil { - t.Fatalf("sending pause command: %v", err) - } - time.Sleep(10 * time.Second) - g.Go(func() error { - return ingestRandomData(ctx, cli[1], "testidx2", "testfield2", 10000) - }) - if err := pauseNode(t, "pilosa3"); err != nil { - t.Fatalf("sending pause command: %v", err) - } - if err := pauseNode(t, "pilosa1"); err != nil { - t.Fatalf("sending pause command: %v", err) - } - time.Sleep(6 * time.Second) - if err := unpauseNode(t, "pilosa3"); err != nil { - t.Fatalf("sending pause command: %v", err) - } - time.Sleep(6 * time.Second) - if err := pauseNode(t, "pilosa2"); err != nil { - t.Fatalf("sending pause command: %v", err) - } - time.Sleep(6 * time.Second) - if err := unpauseNode(t, "pilosa1"); err != nil { - t.Fatalf("sending pause command: %v", err) - } - if err := unpauseNode(t, "pilosa2"); err != nil { - t.Fatalf("sending pause command: %v", err) - } - if err = g.Wait(); err != nil { - t.Fatal(err) - } - waitForStatus(t, cli[0].Status, string(disco.ClusterStateNormal), 30, time.Second, ctx) - - // check data in all three nodes. - for i, c := range cli { - r, err := c.Query(ctx, "testidx1", &pilosa.QueryRequest{Index: "testidx1", Query: "Count(Row(testfield1 = 0))"}) - if err != nil { - t.Fatalf("count querying pilosa%d, %v", i, err) - } - if r.Results[0].(uint64) != 100000 { - t.Fatalf("count on pilosa%d after import is %d", i, r.Results[0].(uint64)) - } - r, err = c.Query(ctx, "testidx2", &pilosa.QueryRequest{Index: "testidx2", Query: "Count(Row(testfield2 = 0))"}) - if err != nil { - t.Fatalf("count querying pilosa%d, %v", i, err) - } - if r.Results[0].(uint64) != 10000 { - t.Fatalf("count on pilosa%d after import is %d", i, r.Results[0].(uint64)) - } - } -} - func waitForStatus(t *testing.T, stator func(context.Context) (string, error), status string, n int, sleep time.Duration, ctx context.Context) { t.Helper() diff --git a/internal/clustertests/docker-compose.yml b/internal/clustertests/docker-compose.yml index e04c4b657..54ac7b22d 100644 --- a/internal/clustertests/docker-compose.yml +++ b/internal/clustertests/docker-compose.yml @@ -17,8 +17,7 @@ services: - pilosanet volumes: - ./results:/results - command: - - "cd /go/src/github.com/featurebasedb/featurebase/cmd/featurebase && /featurebase -test.run=TestRunMain -test.coverprofile=/results/coverage-server1.out server --bind pilosa1:10101 ${CLUSTERTESTS_FB_ARGS}" + command: /featurebase -test.run=TestRunMain -test.coverprofile=/results/coverage-server1.out server --bind pilosa1:10101 ${CLUSTERTESTS_FB_ARGS} pilosa2: build: context: ../.. @@ -36,8 +35,7 @@ services: - pilosanet volumes: - ./results:/results - command: - - "cd /go/src/github.com/featurebasedb/featurebase/cmd/featurebase && /featurebase -test.run=TestRunMain -test.coverprofile=/results/coverage-server2.out server --bind pilosa2:10101 ${CLUSTERTESTS_FB_ARGS}" + command: /featurebase -test.run=TestRunMain -test.coverprofile=/results/coverage-server2.out server --bind pilosa2:10101 ${CLUSTERTESTS_FB_ARGS} pilosa3: build: context: ../.. @@ -55,8 +53,7 @@ services: - pilosanet volumes: - ./results:/results - command: - - "cd /go/src/github.com/featurebasedb/featurebase/cmd/featurebase && /featurebase -test.run=TestRunMain -test.coverprofile=/results/coverage-server3.out server --bind pilosa3:10101 ${CLUSTERTESTS_FB_ARGS}" + command: /featurebase -test.run=TestRunMain -test.coverprofile=/results/coverage-server3.out server --bind pilosa3:10101 ${CLUSTERTESTS_FB_ARGS} client1: build: context: ../.. @@ -76,8 +73,7 @@ services: volumes: - /var/run/docker.sock:/var/run/docker.sock - ./results:/results - command: - - "cd /go/src/github.com/featurebasedb/featurebase/ && go test -mod=vendor -v -count=1 -covermode=atomic -coverprofile=/results/coverage-clustertests.out -coverpkg=./... -json github.com/featurebasedb/featurebase/v3/internal/clustertests | tee /results/report-clustertests.out" + command: "go test -mod=vendor -v -count=1 -covermode=atomic -coverprofile=/results/coverage-clustertests.out -coverpkg=./... -json ./internal/clustertests" fakeidp: build: context: . diff --git a/internal_client.go b/internal_client.go index fce0d2dc8..2b5101567 100644 --- a/internal_client.go +++ b/internal_client.go @@ -13,7 +13,6 @@ import ( "net/http" "net/url" "os" - "path" "sort" "strconv" "strings" @@ -48,6 +47,12 @@ type InternalClient struct { // secret Key for auth across nodes secretKey string + + // pathPrefix is prepended to every URL path. This is used, for example, + // when running a compute nodes as a sub-service of the featurebase command. + // In that case, a path might look like `localhost:8080/compute/schema`, + // where `/compute` is the pathPrefix. + pathPrefix string } // NewInternalClient returns a new instance of InternalClient to connect to host. @@ -113,6 +118,13 @@ func WithClientLogger(log logger.Logger) InternalClientOption { } } +// WithPathPrefix sets the http path prefix. +func WithPathPrefix(prefix string) InternalClientOption { + return func(c *InternalClient) { + c.pathPrefix = prefix + } +} + func noRetryPolicy(ctx context.Context, resp *http.Response, err error) (bool, error) { return false, nil } @@ -257,6 +269,15 @@ func AddAuthToken(ctx context.Context, header *http.Header) { header.Set(authn.RefreshHeaderName, refresh) } +// prefix is a helper function which allows us to provide a pathPrefix value as +// "compute" instead of "/compute". +func (c *InternalClient) prefix() string { + if c.pathPrefix == "" { + return "" + } + return "/" + c.pathPrefix +} + // MaxShardByIndex returns the number of shards on a server by index. func (c *InternalClient) MaxShardByIndex(ctx context.Context) (map[string]uint64, error) { span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.MaxShardByIndex") @@ -267,7 +288,8 @@ func (c *InternalClient) MaxShardByIndex(ctx context.Context) (map[string]uint64 // maxShardByIndex returns the number of shards on a server by index. func (c *InternalClient) maxShardByIndex(ctx context.Context) (map[string]uint64, error) { // Execute request against the host. - u := uriPathToURL(c.defaultURI, "/internal/shards/max") + path := fmt.Sprintf("%s/internal/shards/max", c.prefix()) + u := uriPathToURL(c.defaultURI, path) // Build request. req, err := http.NewRequest("GET", u.String(), nil) @@ -300,7 +322,8 @@ func (c *InternalClient) AvailableShards(ctx context.Context, indexName string) defer span.Finish() // Execute request against the host. - u := uriPathToURL(c.defaultURI, path.Join("/internal/index", indexName, "/shards")) + path := fmt.Sprintf("%s/internal/index/%s/shards", c.prefix(), indexName) + u := uriPathToURL(c.defaultURI, path) // Build request. req, err := http.NewRequest("GET", u.String(), nil) @@ -334,7 +357,7 @@ func (c *InternalClient) SchemaNode(ctx context.Context, uri *pnet.URI, views bo // TODO: /?views parameter will be ignored, till we implement schemator! // Execute request against the host. - u := uri.Path(fmt.Sprintf("/schema?views=%v", views)) + u := uri.Path(fmt.Sprintf("%s/schema?views=%v", c.prefix(), views)) // Build request. req, err := http.NewRequest("GET", u, nil) @@ -366,7 +389,7 @@ func (c *InternalClient) Schema(ctx context.Context) ([]*IndexInfo, error) { defer span.Finish() // Execute request against the host. - u := c.defaultURI.Path("/schema") + u := c.defaultURI.Path(fmt.Sprintf("%s/schema", c.prefix())) // Build request. req, err := http.NewRequest("GET", u, nil) @@ -401,7 +424,7 @@ func (c *InternalClient) MutexCheck(ctx context.Context, uri *pnet.URI, indexNam } // This is not actually a "Path", but reworking this to support queries // is messier than I have resources to pursue just now. - u := uri.Path(fmt.Sprintf("/internal/index/%s/field/%s/mutex-check?details=%t&limit=%d", indexName, fieldName, details, limit)) + u := uri.Path(fmt.Sprintf("%s/internal/index/%s/field/%s/mutex-check?details=%t&limit=%d", c.prefix(), indexName, fieldName, details, limit)) req, err := http.NewRequest("GET", u, nil) if err != nil { return nil, errors.Wrap(err, "creating request") @@ -425,7 +448,7 @@ func (c *InternalClient) MutexCheck(ctx context.Context, uri *pnet.URI, indexNam } func (c *InternalClient) PostSchema(ctx context.Context, uri *pnet.URI, s *Schema, remote bool) error { - u := uri.Path(fmt.Sprintf("/schema?remote=%v", remote)) + u := uri.Path(fmt.Sprintf("%s/schema?remote=%v", c.prefix(), remote)) buf, err := json.Marshal(s) if err != nil { return errors.Wrap(err, "marshalling schema") @@ -477,7 +500,7 @@ func (c *InternalClient) CreateIndex(ctx context.Context, index string, opt Inde } // Create URL & HTTP request. - u := uriPathToURL(&coord.URI, fmt.Sprintf("/index/%s", index)) + u := uriPathToURL(&coord.URI, fmt.Sprintf("%s/index/%s", c.prefix(), index)) req, err := http.NewRequest("POST", u.String(), bytes.NewReader(buf)) if err != nil { return errors.Wrap(err, "creating request") @@ -505,7 +528,7 @@ func (c *InternalClient) FragmentNodes(ctx context.Context, index string, shard defer span.Finish() // Execute request against the host. - u := uriPathToURL(c.defaultURI, "/internal/fragment/nodes") + u := uriPathToURL(c.defaultURI, fmt.Sprintf("%s/internal/fragment/nodes", c.prefix())) u.RawQuery = (url.Values{"index": {index}, "shard": {strconv.FormatUint(shard, 10)}}).Encode() // Build request. @@ -538,7 +561,7 @@ func (c *InternalClient) Nodes(ctx context.Context) ([]*disco.Node, error) { defer span.Finish() // Execute request against the host. - u := uriPathToURL(c.defaultURI, "/internal/nodes") + u := uriPathToURL(c.defaultURI, fmt.Sprintf("%s/internal/nodes", c.prefix())) // Build request. req, err := http.NewRequest("GET", u.String(), nil) @@ -587,7 +610,7 @@ func (c *InternalClient) QueryNode(ctx context.Context, uri *pnet.URI, index str } // Create HTTP request. - u := uri.Path(fmt.Sprintf("/index/%s/query", index)) + u := uri.Path(fmt.Sprintf("%s/index/%s/query", c.prefix(), index)) req, err := http.NewRequest("POST", u, bytes.NewReader(buf)) if err != nil { return nil, errors.Wrap(err, "creating request") @@ -666,7 +689,7 @@ func (c *InternalClient) importNode(ctx context.Context, node *disco.Node, index defer span.Finish() // Create URL & HTTP request. - path := fmt.Sprintf("/index/%s/field/%s/import", index, field) + path := fmt.Sprintf("%s/index/%s/field/%s/import", c.prefix(), index, field) u := nodePathToURL(node, path) vals := url.Values{} @@ -891,7 +914,7 @@ func (c *InternalClient) ImportRoaring(ctx context.Context, uri *pnet.URI, index vals := url.Values{} vals.Set("remote", strconv.FormatBool(remote)) - url := fmt.Sprintf("%s/index/%s/field/%s/import-roaring/%d?%s", uri, index, field, shard, vals.Encode()) + url := fmt.Sprintf("%s%s/index/%s/field/%s/import-roaring/%d?%s", uri, c.prefix(), index, field, shard, vals.Encode()) // Marshal data to protobuf. data, err := c.serializer.Marshal(req) @@ -969,7 +992,7 @@ func (c *InternalClient) exportNodeCSV(ctx context.Context, node *disco.Node, in defer span.Finish() // Create URL. - u := nodePathToURL(node, "/export") + u := nodePathToURL(node, fmt.Sprintf("%s/export", c.prefix())) u.RawQuery = url.Values{ "index": {index}, "field": {field}, @@ -1011,7 +1034,7 @@ func (c *InternalClient) RetrieveShardFromURI(ctx context.Context, index, field, URI: uri, } - u := nodePathToURL(node, "/internal/fragment/data") + u := nodePathToURL(node, fmt.Sprintf("%s/internal/fragment/data", c.prefix())) u.RawQuery = url.Values{ "index": {index}, "field": {field}, @@ -1109,7 +1132,7 @@ func (c *InternalClient) CreateFieldWithOptions(ctx context.Context, index, fiel } // Create URL & HTTP request. - u := uriPathToURL(&coord.URI, fmt.Sprintf("/index/%s/field/%s", index, field)) + u := uriPathToURL(&coord.URI, fmt.Sprintf("%s/index/%s/field/%s", c.prefix(), index, field)) req, err := http.NewRequest("POST", u.String(), bytes.NewReader(buf)) if err != nil { return errors.Wrap(err, "creating request") @@ -1137,7 +1160,7 @@ func (c *InternalClient) SendMessage(ctx context.Context, uri *pnet.URI, msg []b span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.SendMessage") defer span.Finish() - u := uriPathToURL(uri, "/internal/cluster/message") + u := uriPathToURL(uri, fmt.Sprintf("%s/internal/cluster/message", c.prefix())) req, err := http.NewRequest("POST", u.String(), bytes.NewReader(msg)) if err != nil { return errors.Wrap(err, "making new request") @@ -1182,7 +1205,7 @@ func (c *InternalClient) TranslateKeysNode(ctx context.Context, uri *pnet.URI, i } // Create HTTP request. - u := uri.Path("/internal/translate/keys") + u := uri.Path(fmt.Sprintf("%s/internal/translate/keys", c.prefix())) req, err := http.NewRequest("POST", u, bytes.NewReader(buf)) if err != nil { return nil, errors.Wrap(err, "creating request") @@ -1237,7 +1260,7 @@ func (c *InternalClient) TranslateIDsNode(ctx context.Context, uri *pnet.URI, in } // Create HTTP request. - u := uri.Path("/internal/translate/ids") + u := uri.Path(fmt.Sprintf("%s/internal/translate/ids", c.prefix())) req, err := http.NewRequest("POST", u, bytes.NewReader(buf)) if err != nil { return nil, errors.Wrap(err, "creating request") @@ -1272,7 +1295,7 @@ func (c *InternalClient) TranslateIDsNode(ctx context.Context, uri *pnet.URI, in // GetPastQueries retrieves the query history log for the specified node. func (c *InternalClient) GetPastQueries(ctx context.Context, uri *pnet.URI) ([]PastQueryStatus, error) { - u := uri.Path("/query-history?remote=true") + u := uri.Path(fmt.Sprintf("%s/query-history?remote=true", c.prefix())) req, err := http.NewRequest("GET", u, nil) if err != nil { return nil, errors.Wrap(err, "creating request") @@ -1307,7 +1330,7 @@ func (c *InternalClient) FindIndexKeysNode(ctx context.Context, uri *pnet.URI, i defer span.Finish() // Create HTTP request. - u := uriPathToURL(uri, fmt.Sprintf("/internal/translate/index/%s/keys/find", index)) + u := uriPathToURL(uri, fmt.Sprintf("%s/internal/translate/index/%s/keys/find", c.prefix(), index)) reqData, err := json.Marshal(keys) if err != nil { return nil, errors.Wrap(err, "marshalling request") @@ -1357,7 +1380,7 @@ func (c *InternalClient) FindFieldKeysNode(ctx context.Context, uri *pnet.URI, i defer span.Finish() // Create HTTP request. - u := uriPathToURL(uri, fmt.Sprintf("/internal/translate/field/%s/%s/keys/find", index, field)) + u := uriPathToURL(uri, fmt.Sprintf("%s/internal/translate/field/%s/%s/keys/find", c.prefix(), index, field)) q := u.Query() q.Add("remote", "true") u.RawQuery = q.Encode() @@ -1406,7 +1429,7 @@ func (c *InternalClient) CreateIndexKeysNode(ctx context.Context, uri *pnet.URI, defer span.Finish() // Create HTTP request. - u := uriPathToURL(uri, fmt.Sprintf("/internal/translate/index/%s/keys/create", index)) + u := uriPathToURL(uri, fmt.Sprintf("%s/internal/translate/index/%s/keys/create", c.prefix(), index)) reqData, err := json.Marshal(keys) if err != nil { return nil, errors.Wrap(err, "marshalling request") @@ -1456,7 +1479,7 @@ func (c *InternalClient) CreateFieldKeysNode(ctx context.Context, uri *pnet.URI, defer span.Finish() // Create HTTP request. - u := uriPathToURL(uri, fmt.Sprintf("/internal/translate/field/%s/%s/keys/create", index, field)) + u := uriPathToURL(uri, fmt.Sprintf("%s/internal/translate/field/%s/%s/keys/create", c.prefix(), index, field)) q := u.Query() q.Add("remote", "true") u.RawQuery = q.Encode() @@ -1509,7 +1532,7 @@ func (c *InternalClient) MatchFieldKeysNode(ctx context.Context, uri *pnet.URI, defer span.Finish() // Create HTTP request. - u := uriPathToURL(uri, fmt.Sprintf("/internal/translate/field/%s/%s/keys/like", index, field)) + u := uriPathToURL(uri, fmt.Sprintf("%s/internal/translate/field/%s/%s/keys/like", c.prefix(), index, field)) req, err := http.NewRequest("POST", u.String(), strings.NewReader(like)) if err != nil { return nil, errors.Wrap(err, "creating request") @@ -1552,7 +1575,7 @@ func (c *InternalClient) Transactions(ctx context.Context) (map[string]*Transact span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.Transactions") defer span.Finish() - u := uriPathToURL(c.defaultURI, "/transactions") + u := uriPathToURL(c.defaultURI, fmt.Sprintf("%s/transactions", c.prefix())) req, err := http.NewRequest("GET", u.String(), nil) if err != nil { return nil, errors.Wrap(err, "creating transactions request") @@ -1589,7 +1612,7 @@ func (c *InternalClient) StartTransaction(ctx context.Context, id string, timeou // tests, and we want to test requests against all hosts. A robust // client implementation would ensure that these requests go to // the primary. - u := uriPathToURL(c.defaultURI, "/transaction/"+id) + u := uriPathToURL(c.defaultURI, fmt.Sprintf("%s/transaction/%s", c.prefix(), id)) req, err := http.NewRequest("POST", u.String(), bytes.NewReader(buf)) if err != nil { return nil, errors.Wrap(err, "creating post transaction request") @@ -1625,7 +1648,7 @@ func (c *InternalClient) FinishTransaction(ctx context.Context, id string) (*Tra span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.FinishTransaction") defer span.Finish() - u := uriPathToURL(c.defaultURI, "/transaction/"+id+"/finish") + u := uriPathToURL(c.defaultURI, fmt.Sprintf("%s/transaction/%s/finish", c.prefix(), id)) req, err := http.NewRequest("POST", u.String(), nil) if err != nil { return nil, errors.Wrap(err, "creating finish transaction request") @@ -1663,7 +1686,7 @@ func (c *InternalClient) GetTransaction(ctx context.Context, id string) (*Transa // tests, and we want to test requests against all hosts. A robust // client implementation would ensure that these requests go to // the primary. - u := uriPathToURL(c.defaultURI, "/transaction/"+id) + u := uriPathToURL(c.defaultURI, fmt.Sprintf("%s/transaction/%s", c.prefix(), id)) req, err := http.NewRequest("GET", u.String(), nil) if err != nil { return nil, errors.Wrap(err, "creating get transaction request") @@ -2035,7 +2058,7 @@ func (c *InternalClient) RetrieveTranslatePartitionFromURI(ctx context.Context, URI: uri, } - u := nodePathToURL(node, "/internal/translate/data") + u := nodePathToURL(node, fmt.Sprintf("%s/internal/translate/data", c.prefix())) u.RawQuery = url.Values{ "index": {index}, "partition": {strconv.FormatInt(int64(partition), 10)}, @@ -2075,7 +2098,7 @@ func (c *InternalClient) ImportIndexKeys(ctx context.Context, uri *pnet.URI, ind vals := url.Values{} vals.Set("remote", strconv.FormatBool(remote)) - url := fmt.Sprintf("%s/internal/translate/index/%s/%d", uri, index, partitionID) + url := fmt.Sprintf("%s%s/internal/translate/index/%s/%d", uri, c.prefix(), index, partitionID) // Generate HTTP request. httpReq, err := retryablehttp.NewRequest("POST", url, readerFunc) @@ -2109,7 +2132,7 @@ func (c *InternalClient) ImportFieldKeys(ctx context.Context, uri *pnet.URI, ind vals := url.Values{} vals.Set("remote", strconv.FormatBool(remote)) - url := fmt.Sprintf("%s/internal/translate/field/%s/%s", uri, index, field) + url := fmt.Sprintf("%s%s/internal/translate/field/%s/%s", uri, c.prefix(), index, field) // Generate HTTP request. httpReq, err := retryablehttp.NewRequest("POST", url, readerFunc) @@ -2135,7 +2158,7 @@ func (c *InternalClient) ShardReader(ctx context.Context, index string, shard ui defer span.Finish() // Execute request against the host. - u := fmt.Sprintf("%s/internal/index/%s/shard/%d/snapshot", c.defaultURI, index, shard) + u := fmt.Sprintf("%s%s/internal/index/%s/shard/%d/snapshot", c.defaultURI, c.prefix(), index, shard) // Build request. req, err := http.NewRequest("GET", u, nil) @@ -2161,7 +2184,8 @@ func (c *InternalClient) IDAllocDataReader(ctx context.Context) (io.ReadCloser, defer span.Finish() // Build request. - req, err := http.NewRequest("GET", c.defaultURI.String()+"/internal/idalloc/data", nil) + uri := fmt.Sprintf("%s%s/internal/idalloc/data", c.defaultURI, c.prefix()) + req, err := http.NewRequest("GET", uri, nil) if err != nil { return nil, errors.Wrap(err, "creating request") } @@ -2182,7 +2206,7 @@ func (c *InternalClient) IDAllocDataWriter(ctx context.Context, f io.Reader, pri span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.IDAllocDataWriter") defer span.Finish() - u := primary.URI.Path("/internal/idalloc/restore") + u := primary.URI.Path(fmt.Sprintf("%s/internal/idalloc/restore", c.prefix())) // Build request. req, err := http.NewRequest("POST", u, f) @@ -2209,7 +2233,7 @@ func (c *InternalClient) IndexTranslateDataReader(ctx context.Context, index str defer span.Finish() // Execute request against the host. - u := fmt.Sprintf("%s/internal/translate/data?index=%s&partition=%d", c.defaultURI, url.QueryEscape(index), partitionID) + u := fmt.Sprintf("%s%s/internal/translate/data?index=%s&partition=%d", c.defaultURI, c.prefix(), url.QueryEscape(index), partitionID) // Build request. req, err := http.NewRequest("GET", u, nil) @@ -2239,7 +2263,7 @@ func (c *InternalClient) FieldTranslateDataReader(ctx context.Context, index, fi defer span.Finish() // Execute request against the host. - u := fmt.Sprintf("%s/internal/translate/data?index=%s&field=%s", c.defaultURI, url.QueryEscape(index), url.QueryEscape(field)) + u := fmt.Sprintf("%s%s/internal/translate/data?index=%s&field=%s", c.defaultURI, c.prefix(), url.QueryEscape(index), url.QueryEscape(field)) // Build request. req, err := http.NewRequest("GET", u, nil) @@ -2268,7 +2292,7 @@ func (c *InternalClient) Status(ctx context.Context) (string, error) { defer span.Finish() // Execute request against the host. - u := c.defaultURI.Path("/status") + u := c.defaultURI.Path(fmt.Sprintf("%s/status", c.prefix())) // Build request. req, err := http.NewRequest("GET", u, nil) @@ -2299,7 +2323,7 @@ func (c *InternalClient) PartitionNodes(ctx context.Context, partitionID int) ([ defer span.Finish() // Execute request against the host. - u := uriPathToURL(c.defaultURI, "/internal/partition/nodes") + u := uriPathToURL(c.defaultURI, fmt.Sprintf("%s/internal/partition/nodes", c.prefix())) u.RawQuery = (url.Values{"partition": {strconv.FormatInt(int64(partitionID), 10)}}).Encode() // Build request. @@ -2331,7 +2355,7 @@ func (c *InternalClient) SetInternalAPI(api *API) { } func (c *InternalClient) OAuthConfig() (rsp oauth2.Config, err error) { - u := uriPathToURL(c.defaultURI, "/internal/oauth-config") + u := uriPathToURL(c.defaultURI, fmt.Sprintf("%s/internal/oauth-config", c.prefix())) req, err := http.NewRequest("GET", u.String(), nil) if err != nil { diff --git a/logger/logger.go b/logger/logger.go index 6ef8cc1f8..60b0fce3c 100644 --- a/logger/logger.go +++ b/logger/logger.go @@ -7,6 +7,7 @@ import ( "fmt" "io" "log" + "os" "sync" "time" @@ -40,6 +41,8 @@ func LevelPrefix(level int) string { return [...]string{"PANIC: ", "ERROR: ", "WARN: ", "INFO: ", "DEBUG: "}[level] } +var StderrLogger = NewStandardLogger(os.Stderr) + // NopLogger represents a Logger that doesn't do anything. var NopLogger Logger = &nopLogger{} diff --git a/rbf/cfg/cfg.go b/rbf/cfg/cfg.go index 3c5134b70..3fd593bd6 100644 --- a/rbf/cfg/cfg.go +++ b/rbf/cfg/cfg.go @@ -64,15 +64,23 @@ func NewDefaultConfig() *Config { } } -func (cfg *Config) DefineFlags(flags *pflag.FlagSet) { +func (cfg *Config) DefineFlags(flags *pflag.FlagSet, prefix string) { + // pre applies prefix to s when a prefix is provided. + pre := func(s string) string { + if prefix == "" { + return s + } + return prefix + "." + s + } + default0 := NewDefaultConfig() - flags.Int64Var(&cfg.MaxSize, "rbf.max-db-size", default0.MaxSize, "RBF maximum size in bytes of a database file (distinct from a WAL file)") - flags.Int64Var(&cfg.MaxWALSize, "rbf.max-wal-size", default0.MaxWALSize, "RBF maximum size in bytes of a WAL file (distinct from a DB file)") - flags.Int64Var(&cfg.MinWALCheckpointSize, "rbf.min-wal-checkpoint-size", default0.MinWALCheckpointSize, "RBF minimum size in bytes of a WAL file before attempting checkpoint") - flags.Int64Var(&cfg.MaxWALCheckpointSize, "rbf.max-wal-checkpoint-size", default0.MaxWALCheckpointSize, "RBF maximum size in bytes of a WAL file before forcing checkpoint") + flags.Int64Var(&cfg.MaxSize, pre("rbf.max-db-size"), default0.MaxSize, "RBF maximum size in bytes of a database file (distinct from a WAL file)") + flags.Int64Var(&cfg.MaxWALSize, pre("rbf.max-wal-size"), default0.MaxWALSize, "RBF maximum size in bytes of a WAL file (distinct from a DB file)") + flags.Int64Var(&cfg.MinWALCheckpointSize, pre("rbf.min-wal-checkpoint-size"), default0.MinWALCheckpointSize, "RBF minimum size in bytes of a WAL file before attempting checkpoint") + flags.Int64Var(&cfg.MaxWALCheckpointSize, pre("rbf.max-wal-checkpoint-size"), default0.MaxWALCheckpointSize, "RBF maximum size in bytes of a WAL file before forcing checkpoint") // renamed from --rbf-fsync to just --fsync because now it applies to all Tx backends. - flags.BoolVar(&cfg.FsyncEnabled, "fsync", default0.FsyncEnabled, "enable fsync fully safe flush-to-disk") - flags.BoolVar(&cfg.FsyncWALEnabled, "fsync-wal", default0.FsyncWALEnabled, "enable fsync on write-ahead log") - flags.Int64Var(&cfg.CursorCacheSize, "rbf.cursor-cache-size", default0.CursorCacheSize, "how big a Cursor arena to maintain. 0 means use sync.Pool with dynamic sizing. Note that <= 20 is needed to pass CI. Controls the memory footprint of rbf.") + flags.BoolVar(&cfg.FsyncEnabled, pre("fsync"), default0.FsyncEnabled, "enable fsync fully safe flush-to-disk") + flags.BoolVar(&cfg.FsyncWALEnabled, pre("fsync-wal"), default0.FsyncWALEnabled, "enable fsync on write-ahead log") + flags.Int64Var(&cfg.CursorCacheSize, pre("rbf.cursor-cache-size"), default0.CursorCacheSize, "how big a Cursor arena to maintain. 0 means use sync.Pool with dynamic sizing. Note that <= 20 is needed to pass CI. Controls the memory footprint of rbf.") } diff --git a/row.go b/row.go index 941bc67eb..cefa862ae 100644 --- a/row.go +++ b/row.go @@ -13,7 +13,7 @@ import ( // Row is a set of integers (the associated columns). type Row struct { - segments []rowSegment + Segments []RowSegment // String keys translated to/from segment columns. Keys []string @@ -57,8 +57,8 @@ func (r *Row) Clone() (clone *Row) { Field: r.Field, } - for _, seg := range r.segments { - segClone := rowSegment{ + for _, seg := range r.Segments { + segClone := RowSegment{ shard: seg.shard, writable: true, // we know it is safe; it is a copy. n: seg.n, @@ -67,7 +67,7 @@ func (r *Row) Clone() (clone *Row) { segClone.data = seg.data.Clone() // *roaring.Bitmap } //segClone.InvalidateCount() // not needed? - clone.segments = append(clone.segments, segClone) + clone.Segments = append(clone.Segments, segClone) } return clone } @@ -83,13 +83,13 @@ func NewRowFromBitmap(b *roaring.Bitmap) *Row { rowNum := uint64(0) for col, ok := b.MinAt(rowNum * ShardWidth); ok; col, ok = b.MinAt(rowNum * ShardWidth) { rowNum = col / ShardWidth - seg := rowSegment{ + seg := RowSegment{ shard: rowNum, data: b.OffsetRange(rowNum*ShardWidth, rowNum*ShardWidth, (rowNum+1)*ShardWidth), writable: true, } seg.n = seg.data.Count() - r.segments = append(r.segments, seg) + r.Segments = append(r.Segments, seg) rowNum++ } return r @@ -99,15 +99,15 @@ func NewRowFromBitmap(b *roaring.Bitmap) *Row { // bitmaps and rowSegments based on shard width. func NewRowFromRoaring(data []byte) *Row { bitmaps, shards := roaring.RoaringToBitmaps(data, ShardWidth) - r := &Row{segments: make([]rowSegment, len(bitmaps))} + r := &Row{Segments: make([]RowSegment, len(bitmaps))} for i := range bitmaps { - segment := rowSegment{ + segment := RowSegment{ shard: shards[i], data: bitmaps[i], writable: false, n: bitmaps[i].Count(), } - r.segments[i] = segment + r.Segments[i] = segment } return r } @@ -126,8 +126,8 @@ func (r *Row) ToTable() (*pb.TableResponse, error) { // Hash calculate checksum code be useful in block hash join func (r *Row) Hash() uint64 { hash := uint64(0) - for i := range r.segments { - hash = r.segments[i].data.Hash(hash) + for i := range r.Segments { + hash = r.Segments[i].data.Hash(hash) } return hash } @@ -170,20 +170,20 @@ func (r *Row) ToRows(callback func(*pb.RowResponse) error) error { // Roaring returns the row treated as a unified roaring bitmap. func (r *Row) Roaring() []byte { - bitmaps := make([]*roaring.Bitmap, len(r.segments)) - for i := range r.segments { - bitmaps[i] = r.segments[i].data + bitmaps := make([]*roaring.Bitmap, len(r.Segments)) + for i := range r.Segments { + bitmaps[i] = r.Segments[i].data } return roaring.BitmapsToRoaring(bitmaps) } // IsEmpty returns true if the row doesn't contain any set bits. func (r *Row) IsEmpty() bool { - if len(r.segments) == 0 { + if len(r.Segments) == 0 { return true } - for i := range r.segments { - if r.segments[i].n > 0 { + for i := range r.Segments { + if r.Segments[i].n > 0 { return false } @@ -192,16 +192,16 @@ func (r *Row) IsEmpty() bool { } func (r *Row) Freeze() { - for _, s := range r.segments { + for _, s := range r.Segments { s.Freeze() } } // Merge merges data from other into r. func (r *Row) Merge(other *Row) { - var segments []rowSegment + var segments []RowSegment - itr := newMergeSegmentIterator(r.segments, other.segments) + itr := newMergeSegmentIterator(r.Segments, other.Segments) for s0, s1 := itr.next(); s0 != nil || s1 != nil; s0, s1 = itr.next() { // Use the other row's data if segment is missing. if s0 == nil { @@ -217,7 +217,7 @@ func (r *Row) Merge(other *Row) { segments = append(segments, *s0) } - r.segments = segments + r.Segments = segments r.invalidateCount() } @@ -225,7 +225,7 @@ func (r *Row) Merge(other *Row) { func (r *Row) intersectionCount(other *Row) uint64 { var n uint64 - itr := newMergeSegmentIterator(r.segments, other.segments) + itr := newMergeSegmentIterator(r.Segments, other.Segments) for s0, s1 := itr.next(); s0 != nil || s1 != nil; s0, s1 = itr.next() { // Ignore non-overlapping segments. if s0 == nil || s1 == nil { @@ -239,9 +239,9 @@ func (r *Row) intersectionCount(other *Row) uint64 { // Intersect returns the itersection of r and other. func (r *Row) Intersect(other *Row) *Row { - var segments []rowSegment + var segments []RowSegment - itr := newMergeSegmentIterator(r.segments, other.segments) + itr := newMergeSegmentIterator(r.Segments, other.Segments) for s0, s1 := itr.next(); s0 != nil || s1 != nil; s0, s1 = itr.next() { // Ignore non-overlapping segments. if s0 == nil || s1 == nil { @@ -250,12 +250,12 @@ func (r *Row) Intersect(other *Row) *Row { segments = append(segments, *s0.Intersect(s1)) } - return &Row{segments: segments} + return &Row{Segments: segments} } // Any returns true if row contains any bits. func (r *Row) Any() bool { - for _, s := range r.segments { + for _, s := range r.Segments { if s.data.Any() { return true } @@ -265,9 +265,9 @@ func (r *Row) Any() bool { // Xor returns the xor of r and other. func (r *Row) Xor(other *Row) *Row { - var segments []rowSegment + var segments []RowSegment - itr := newMergeSegmentIterator(r.segments, other.segments) + itr := newMergeSegmentIterator(r.Segments, other.Segments) for s0, s1 := itr.next(); s0 != nil || s1 != nil; s0, s1 = itr.next() { if s1 == nil { segments = append(segments, *s0) @@ -280,21 +280,21 @@ func (r *Row) Xor(other *Row) *Row { segments = append(segments, *s0.Xor(s1)) } - return &Row{segments: segments} + return &Row{Segments: segments} } // Union returns the bitwise union of r and other. func (r *Row) Union(others ...*Row) *Row { - segments := make([][]rowSegment, 0, len(others)+1) - if len(r.segments) > 0 { - segments = append(segments, r.segments) + segments := make([][]RowSegment, 0, len(others)+1) + if len(r.Segments) > 0 { + segments = append(segments, r.Segments) } - nextSegs := make([][]rowSegment, 0, len(others)+1) - toProcess := make([]*rowSegment, 0, len(others)+1) - var output []rowSegment + nextSegs := make([][]RowSegment, 0, len(others)+1) + toProcess := make([]*RowSegment, 0, len(others)+1) + var output []RowSegment for _, other := range others { - if len(other.segments) > 0 { - segments = append(segments, other.segments) + if len(other.Segments) > 0 { + segments = append(segments, other.Segments) } } for len(segments) > 0 { @@ -325,21 +325,21 @@ func (r *Row) Union(others ...*Row) *Row { output = append(output, *toProcess[0].Union(toProcess[1:]...)) } } - return &Row{Index: r.Index, Field: r.Field, segments: output} + return &Row{Index: r.Index, Field: r.Field, Segments: output} } // Difference returns the diff of r and other. func (r *Row) Difference(others ...*Row) *Row { - var output []rowSegment - o := make(map[uint64][]*rowSegment) + var output []RowSegment + o := make(map[uint64][]*RowSegment) for x := range others { - for y := range others[x].segments { - segment := others[x].segments[y] + for y := range others[x].Segments { + segment := others[x].Segments[y] o[segment.shard] = append(o[segment.shard], &segment) } } - for _, segment := range r.segments { + for _, segment := range r.Segments { dest, ok := o[segment.shard] if ok { @@ -348,7 +348,7 @@ func (r *Row) Difference(others ...*Row) *Row { output = append(output, segment) } } - return &Row{segments: output} + return &Row{Segments: output} } // Shift returns the bitwise shift of r by n bits. @@ -378,17 +378,17 @@ func (r *Row) Shift(n int64) (*Row, error) { } work := r - var segments []rowSegment + var segments []RowSegment for i := int64(0); i < n; i++ { segments = segments[:0] - for _, segment := range work.segments { + for _, segment := range work.Segments { shifted, err := segment.Shift() if err != nil { return nil, errors.Wrap(err, "shifting row segment") } segments = append(segments, *shifted) } - work = &Row{segments: segments} + work = &Row{Segments: segments} } return work, nil @@ -399,50 +399,45 @@ func (r *Row) SetBit(i uint64) (changed bool) { return r.createSegmentIfNotExists(i / ShardWidth).SetBit(i) } -// Segments returns a list of all segments in the row. -func (r *Row) Segments() []rowSegment { - return r.segments -} - // segment returns a segment for a given shard. // Returns nil if segment does not exist. -func (r *Row) segment(shard uint64) *rowSegment { - if i := sort.Search(len(r.segments), func(i int) bool { - return r.segments[i].shard >= shard - }); i < len(r.segments) && r.segments[i].shard == shard { - return &r.segments[i] +func (r *Row) segment(shard uint64) *RowSegment { + if i := sort.Search(len(r.Segments), func(i int) bool { + return r.Segments[i].shard >= shard + }); i < len(r.Segments) && r.Segments[i].shard == shard { + return &r.Segments[i] } return nil } -func (r *Row) createSegmentIfNotExists(shard uint64) *rowSegment { - i := sort.Search(len(r.segments), func(i int) bool { - return r.segments[i].shard >= shard +func (r *Row) createSegmentIfNotExists(shard uint64) *RowSegment { + i := sort.Search(len(r.Segments), func(i int) bool { + return r.Segments[i].shard >= shard }) // Return exact match. - if i < len(r.segments) && r.segments[i].shard == shard { - return &r.segments[i] + if i < len(r.Segments) && r.Segments[i].shard == shard { + return &r.Segments[i] } // Insert new segment. - r.segments = append(r.segments, rowSegment{data: roaring.NewSliceBitmap()}) - if i < len(r.segments) { - copy(r.segments[i+1:], r.segments[i:]) + r.Segments = append(r.Segments, RowSegment{data: roaring.NewSliceBitmap()}) + if i < len(r.Segments) { + copy(r.Segments[i+1:], r.Segments[i:]) } - r.segments[i] = rowSegment{ + r.Segments[i] = RowSegment{ data: roaring.NewSliceBitmap(), shard: shard, writable: true, } - return &r.segments[i] + return &r.Segments[i] } // invalidateCount updates the cached count in the row. func (r *Row) invalidateCount() { - for i := range r.segments { - r.segments[i].InvalidateCount() + for i := range r.Segments { + r.Segments[i].InvalidateCount() } } @@ -453,8 +448,8 @@ func (r *Row) Count() uint64 { // Count(Distinct()) on an empty field panics here return n } - for i := range r.segments { - n += r.segments[i].Count() + for i := range r.Segments { + n += r.Segments[i].Count() } return n } @@ -479,8 +474,8 @@ func (r *Row) Columns() []uint64 { return nil } a := make([]uint64, 0, r.Count()) - for i := range r.segments { - a = append(a, r.segments[i].Columns()...) + for i := range r.Segments { + a = append(a, r.Segments[i].Columns()...) } return a } @@ -488,18 +483,18 @@ func (r *Row) Columns() []uint64 { // Includes returns true if the row contains the given column. func (r *Row) Includes(col uint64) bool { shard := col / ShardWidth - for i := range r.segments { - if r.segments[i].shard == shard { - return r.segments[i].data.Contains(col) + for i := range r.Segments { + if r.Segments[i].shard == shard { + return r.Segments[i].data.Contains(col) } } return false } -// rowSegment holds a subset of a row. +// RowSegment holds a subset of a row. // This could point to a mmapped roaring bitmap or an in-memory bitmap. The // width of the segment will always match the shard width. -type rowSegment struct { +type RowSegment struct { // Shard this segment belongs to shard uint64 @@ -513,7 +508,11 @@ type rowSegment struct { n uint64 } -func (s *rowSegment) Freeze() { +func (s *RowSegment) Shard() uint64 { + return s.shard +} + +func (s *RowSegment) Freeze() { s.data = s.data.Freeze() } @@ -530,7 +529,7 @@ func (s *rowSegment) Raw() (uint64, []byte) { // Merge adds chunks from other to s. // Chunks in s are overwritten if they exist in other. -func (s *rowSegment) Merge(other *rowSegment) { +func (s *RowSegment) Merge(other *RowSegment) { s.ensureWritable() itr := other.data.Iterator() @@ -540,15 +539,15 @@ func (s *rowSegment) Merge(other *rowSegment) { } // IntersectionCount returns the number of intersections between s and other. -func (s *rowSegment) IntersectionCount(other *rowSegment) uint64 { +func (s *RowSegment) IntersectionCount(other *RowSegment) uint64 { return s.data.IntersectionCount(other.data) } // Intersect returns the itersection of s and other. -func (s *rowSegment) Intersect(other *rowSegment) *rowSegment { +func (s *RowSegment) Intersect(other *RowSegment) *RowSegment { data := s.data.Intersect(other.data) - return &rowSegment{ + return &RowSegment{ data: data, shard: s.shard, n: data.Count(), @@ -556,14 +555,14 @@ func (s *rowSegment) Intersect(other *rowSegment) *rowSegment { } // Union returns the bitwise union of s and other. -func (s *rowSegment) Union(others ...*rowSegment) *rowSegment { +func (s *RowSegment) Union(others ...*RowSegment) *RowSegment { datas := make([]*roaring.Bitmap, len(others)) for i, other := range others { datas[i] = other.data } data := s.data.Union(datas...) - return &rowSegment{ + return &RowSegment{ data: data, shard: s.shard, n: data.Count(), @@ -571,14 +570,14 @@ func (s *rowSegment) Union(others ...*rowSegment) *rowSegment { } // Difference returns the diff of s and other. -func (s *rowSegment) Difference(others ...*rowSegment) *rowSegment { +func (s *RowSegment) Difference(others ...*RowSegment) *RowSegment { datas := make([]*roaring.Bitmap, len(others)) for i, other := range others { datas[i] = other.data } data := s.data.Difference(datas...) - return &rowSegment{ + return &RowSegment{ data: data, shard: s.shard, n: data.Count(), @@ -586,10 +585,10 @@ func (s *rowSegment) Difference(others ...*rowSegment) *rowSegment { } // Xor returns the xor of s and other. -func (s *rowSegment) Xor(other *rowSegment) *rowSegment { +func (s *RowSegment) Xor(other *RowSegment) *RowSegment { data := s.data.Xor(other.data) - return &rowSegment{ + return &RowSegment{ data: data, shard: s.shard, n: data.Count(), @@ -597,7 +596,7 @@ func (s *rowSegment) Xor(other *rowSegment) *rowSegment { } // Shift returns s shifted by 1 bit. -func (s *rowSegment) Shift() (*rowSegment, error) { +func (s *RowSegment) Shift() (*RowSegment, error) { // TODO: deal with overflow // See issue: https://github.com/molecula/pilosa/issues/403 data, err := s.data.Shift(1) @@ -605,7 +604,7 @@ func (s *rowSegment) Shift() (*rowSegment, error) { return nil, errors.Wrap(err, "shifting roaring data") } - return &rowSegment{ + return &RowSegment{ data: data, shard: s.shard, n: data.Count(), @@ -613,7 +612,7 @@ func (s *rowSegment) Shift() (*rowSegment, error) { } // SetBit sets the i-th column of the row. -func (s *rowSegment) SetBit(i uint64) (changed bool) { +func (s *RowSegment) SetBit(i uint64) (changed bool) { s.ensureWritable() changed, _ = s.data.Add(i) if changed { @@ -623,7 +622,7 @@ func (s *rowSegment) SetBit(i uint64) (changed bool) { } // ClearBit clears the i-th column of the row. -func (s *rowSegment) ClearBit(i uint64) (changed bool) { +func (s *RowSegment) ClearBit(i uint64) (changed bool) { s.ensureWritable() changed, _ = s.data.Remove(i) @@ -634,12 +633,12 @@ func (s *rowSegment) ClearBit(i uint64) (changed bool) { } // InvalidateCount updates the cached count in the row. -func (s *rowSegment) InvalidateCount() { +func (s *RowSegment) InvalidateCount() { s.n = s.data.Count() } // Columns returns a list of all columns set in the segment. -func (s *rowSegment) Columns() []uint64 { +func (s *RowSegment) Columns() []uint64 { a := make([]uint64, 0, s.Count()) itr := s.data.Iterator() for v, eof := itr.Next(); !eof; v, eof = itr.Next() { @@ -649,10 +648,10 @@ func (s *rowSegment) Columns() []uint64 { } // Count returns the number of set columns in the row. -func (s *rowSegment) Count() uint64 { return s.n } +func (s *RowSegment) Count() uint64 { return s.n } // ensureWritable clones the segment if it is pointing to non-writable data. -func (s *rowSegment) ensureWritable() { +func (s *RowSegment) ensureWritable() { if s.writable { return } @@ -667,16 +666,16 @@ func (s *rowSegment) ensureWritable() { // mergeSegmentIterator produces an iterator that loops through two sets of segments. type mergeSegmentIterator struct { - a0, a1 []rowSegment + a0, a1 []RowSegment } // newMergeSegmentIterator returns a new instance of mergeSegmentIterator. -func newMergeSegmentIterator(a0, a1 []rowSegment) mergeSegmentIterator { +func newMergeSegmentIterator(a0, a1 []RowSegment) mergeSegmentIterator { return mergeSegmentIterator{a0: a0, a1: a1} } // next returns the next set of segments. -func (itr *mergeSegmentIterator) next() (s0, s1 *rowSegment) { +func (itr *mergeSegmentIterator) next() (s0, s1 *RowSegment) { // Find current segments. if len(itr.a0) > 0 { s0 = &itr.a0[0] diff --git a/server.go b/server.go index 67a250307..5bf93b06b 100644 --- a/server.go +++ b/server.go @@ -17,16 +17,18 @@ import ( uuid "github.com/satori/go.uuid" - "github.com/featurebasedb/featurebase/v3/disco" - "github.com/featurebasedb/featurebase/v3/logger" - pnet "github.com/featurebasedb/featurebase/v3/net" - rbfcfg "github.com/featurebasedb/featurebase/v3/rbf/cfg" - "github.com/featurebasedb/featurebase/v3/roaring" - "github.com/featurebasedb/featurebase/v3/sql3" - "github.com/featurebasedb/featurebase/v3/sql3/parser" - planner_types "github.com/featurebasedb/featurebase/v3/sql3/planner/types" - "github.com/featurebasedb/featurebase/v3/stats" - "github.com/featurebasedb/featurebase/v3/storage" + "github.com/molecula/featurebase/v3/dax/computer" + "github.com/molecula/featurebase/v3/dax/inmem" + "github.com/molecula/featurebase/v3/disco" + "github.com/molecula/featurebase/v3/logger" + pnet "github.com/molecula/featurebase/v3/net" + rbfcfg "github.com/molecula/featurebase/v3/rbf/cfg" + "github.com/molecula/featurebase/v3/roaring" + "github.com/molecula/featurebase/v3/sql3" + "github.com/molecula/featurebase/v3/sql3/parser" + planner_types "github.com/molecula/featurebase/v3/sql3/planner/types" + "github.com/molecula/featurebase/v3/stats" + "github.com/molecula/featurebase/v3/storage" "github.com/pkg/errors" "golang.org/x/sync/errgroup" @@ -92,6 +94,10 @@ type Server struct { // nolint: maligned queryHistoryLength int executionPlannerFn ExecutionPlannerFn + + writeLogReader computer.WriteLogReader + writeLogWriter computer.WriteLogWriter + snapshotReadWriter computer.SnapshotReadWriter } type ExecutionPlannerFn func(executor Executor, api *API, sql string) sql3.CompilePlanner @@ -419,6 +425,15 @@ func OptServerPartitionAssigner(p string) ServerOption { } } +// OptServerWriteLogReader provides an implemenation of the WriteLogReader +// interface. +func OptServerWriteLogReader(wlr computer.WriteLogReader) ServerOption { + return func(s *Server) error { + s.writeLogReader = wlr + return nil + } +} + func OptServerExecutionPlannerFn(fn ExecutionPlannerFn) ServerOption { return func(s *Server) error { s.executionPlannerFn = fn @@ -426,6 +441,32 @@ func OptServerExecutionPlannerFn(fn ExecutionPlannerFn) ServerOption { } } +// OptServerWriteLogWriter provides an implemenation of the WriteLogWriter +// interface. +func OptServerWriteLogWriter(wlw computer.WriteLogWriter) ServerOption { + return func(s *Server) error { + s.writeLogWriter = wlw + return nil + } +} + +// OptServerSnapshotReadWriter provides an implemenation of the +// SnapshotReadWriter interface. +func OptServerSnapshotReadWriter(snap computer.SnapshotReadWriter) ServerOption { + return func(s *Server) error { + s.snapshotReadWriter = snap + return nil + } +} + +// OptServerIsComputeNode specifies that this node is running as a DAX compute node. +func OptServerIsComputeNode(is bool) ServerOption { + return func(s *Server) error { + s.cluster.isComputeNode = is + return nil + } +} + // NewServer returns a new instance of Server. func NewServer(opts ...ServerOption) (*Server, error) { cluster := newCluster() @@ -510,12 +551,23 @@ func NewServer(opts ...ServerOption) (*Server, error) { s.holder.Logger.Infof("cwd: %v", cwd) s.holder.Logger.Infof("cmd line: %v", strings.Join(os.Args, " ")) + // The compute nodes keep a local cache of the VersionStore which applies + // only to the data (shard, partitions, fields) managed by the compute node + // (as opposed to the VersionStore in MDS which keeps information about all + // data). It would be okay for this to be an in-memory implementation as + // long as the compute node isn't expected to survive a restart; in that + // case, it would be necessary to use an implementation which saves state + // somewhere, such as local disk. + versionStore := inmem.NewVersionStore() + s.cluster.Path = path s.cluster.logger = s.logger s.cluster.holder = s.holder s.cluster.disCo = s.disCo s.cluster.noder = s.noder s.cluster.sharder = s.sharder + s.cluster.writeLogWriter = s.writeLogWriter + s.cluster.versionStore = versionStore // Append the NodeID tag to stats. s.holder.Stats = s.holder.Stats.WithTags(fmt.Sprintf("node_id:%s", s.nodeID)) @@ -531,6 +583,7 @@ func NewServer(opts ...ServerOption) (*Server, error) { s.holder.broadcaster = s s.holder.sharder = s.sharder s.holder.serializer = s.serializer + s.holder.versionStore = versionStore // Initial stats must be invoked after the executor obtains reference to the holder. s.executor.InitStats() diff --git a/server/config.go b/server/config.go index 351cf8406..4e59e6400 100644 --- a/server/config.go +++ b/server/config.go @@ -52,6 +52,20 @@ type Config struct { // Name a unique name for this node in the cluster. Name string `toml:"name"` + // MDSAddress is the location at which this node should register itself and + // retrieve its instructions. For example, after registring, the MDS service + // might tell this node that it is responsible for specific shards for a + // particular index. + MDSAddress string `toml:"mds-address"` + + // WriteLogger is the location at which this node should read/write change + // logs. + WriteLogger string `toml:"write-logger"` + + // Snapshotter is the location at which this node should read/write + // snapshots. + Snapshotter string `toml:"snapshotter"` + // DataDir is the directory where Pilosa stores both indexed data and // running state such as cluster topology information. DataDir string `toml:"data-dir"` @@ -62,6 +76,10 @@ type Config struct { // BindGRPC is the host:port on which Pilosa will bind for gRPC. BindGRPC string `toml:"bind-grpc"` + // Listener is an already-bound listener to use for http. + //Listener *net.TCPListener + Listener net.Listener + // GRPCListener is an already-bound listener to use for gRPC. // This is for use by test infrastructure, where it's useful to // be able to dynamically generate the bindings by actually binding @@ -123,6 +141,12 @@ type Config struct { // don't exhaust the goroutine limit. ImportWorkerPoolSize int `toml:"-"` + // DirectiveWorkerPoolSize controls how many goroutines are created for + // processing a Directive (i.e. concurrently loading key/partition/shard + // data from shapshotter and writelogger) on a compute node. Defaults to + // runtime.NumCPU(). + DirectiveWorkerPoolSize int `toml:"-"` + // Limits the total amount of memory to be used by Extract() & SELECT queries. MaxQueryMemory int64 `toml:"max-query-memory"` @@ -183,6 +207,10 @@ type Config struct { EndpointEnabled bool `toml:"endpoint-enabled"` } `toml:"sql"` + // CheckInTimeout is the amount of time between compute node check-ins to + // MDS. + CheckInInterval time.Duration `toml:"check-in-interval"` + // Storage.Backend determines which Tx implementation the holder/Index will // use; one of the available transactional-storage engines. Choices are // listed in the string constants below. Should be one of "roaring" or @@ -337,12 +365,16 @@ func NewConfig() *Config { WorkerPoolSize: runtime.NumCPU(), ImportWorkerPoolSize: runtime.NumCPU(), + DirectiveWorkerPoolSize: runtime.NumCPU(), + Storage: storage.NewDefaultConfig(), RBFConfig: rbfcfg.NewDefaultConfig(), QueryHistoryLength: 100, LongQueryTime: toml.Duration(-time.Minute), + + CheckInInterval: 5 * time.Second, } // Cluster config. diff --git a/server/server.go b/server/server.go index 662205f56..ddab6cfb7 100644 --- a/server/server.go +++ b/server/server.go @@ -15,6 +15,7 @@ import ( "log" "math/rand" "net" + "net/http" "os" "os/signal" "path/filepath" @@ -25,27 +26,30 @@ import ( "syscall" "time" + "github.com/molecula/featurebase/v3/dax" "golang.org/x/sync/errgroup" - pilosa "github.com/featurebasedb/featurebase/v3" - "github.com/featurebasedb/featurebase/v3/authn" - "github.com/featurebasedb/featurebase/v3/authz" - "github.com/featurebasedb/featurebase/v3/batch" - "github.com/featurebasedb/featurebase/v3/boltdb" - "github.com/featurebasedb/featurebase/v3/encoding/proto" - petcd "github.com/featurebasedb/featurebase/v3/etcd" - "github.com/featurebasedb/featurebase/v3/gcnotify" - "github.com/featurebasedb/featurebase/v3/gopsutil" - "github.com/featurebasedb/featurebase/v3/logger" - pnet "github.com/featurebasedb/featurebase/v3/net" - "github.com/featurebasedb/featurebase/v3/prometheus" - "github.com/featurebasedb/featurebase/v3/sql3" - "github.com/featurebasedb/featurebase/v3/sql3/planner" - "github.com/featurebasedb/featurebase/v3/statik" - "github.com/featurebasedb/featurebase/v3/stats" - "github.com/featurebasedb/featurebase/v3/statsd" - "github.com/featurebasedb/featurebase/v3/syswrap" - "github.com/featurebasedb/featurebase/v3/testhook" + pilosa "github.com/molecula/featurebase/v3" + "github.com/molecula/featurebase/v3/authn" + "github.com/molecula/featurebase/v3/authz" + "github.com/molecula/featurebase/v3/batch" + "github.com/molecula/featurebase/v3/boltdb" + "github.com/molecula/featurebase/v3/dax/computer" + "github.com/molecula/featurebase/v3/dax/computer/alpha" + "github.com/molecula/featurebase/v3/encoding/proto" + petcd "github.com/molecula/featurebase/v3/etcd" + "github.com/molecula/featurebase/v3/gcnotify" + "github.com/molecula/featurebase/v3/gopsutil" + "github.com/molecula/featurebase/v3/logger" + pnet "github.com/molecula/featurebase/v3/net" + "github.com/molecula/featurebase/v3/prometheus" + "github.com/molecula/featurebase/v3/sql3" + "github.com/molecula/featurebase/v3/sql3/planner" + "github.com/molecula/featurebase/v3/statik" + "github.com/molecula/featurebase/v3/stats" + "github.com/molecula/featurebase/v3/statsd" + "github.com/molecula/featurebase/v3/syswrap" + "github.com/molecula/featurebase/v3/testhook" "github.com/pelletier/go-toml" "github.com/pkg/errors" ) @@ -75,7 +79,12 @@ type Command struct { logger loggerLogger queryLogger loggerLogger + mds pilosa.MDS + writeLogger pilosa.WriteLogger + snapshotter pilosa.Snapshotter + Handler pilosa.HandlerI + httpHandler http.Handler grpcServer *grpcServer grpcLn net.Listener API *pilosa.API @@ -86,6 +95,10 @@ type Command struct { serverOptions []pilosa.ServerOption auth *authn.Auth + + // isComputeNode is set to true if this node is running as a DAX compute + // node. + isComputeNode bool } type CommandOption func(c *Command) error @@ -111,6 +124,8 @@ func OptCommandConfig(config *Config) CommandOption { c.Config.Etcd = config.Etcd c.Config.Auth = config.Auth c.Config.TLS = config.TLS + c.Config.MDSAddress = config.MDSAddress + c.Config.WriteLogger = config.WriteLogger return nil } c.Config = config @@ -118,6 +133,41 @@ func OptCommandConfig(config *Config) CommandOption { } } +// OptCommandSetConfig was added because OptCommandConfig only sets a small +// sub-set of the config options (it doesn't seem to be used for anything but +// tests). We need a functional option which sets the full Config. +func OptCommandSetConfig(config *Config) CommandOption { + return func(c *Command) error { + defer c.Config.MustValidate() + c.Config = config + return nil + } +} + +// OptCommandInjections injects the interface implementations. +func OptCommandInjections(inj Injections) CommandOption { + return func(c *Command) error { + if inj.MDS != nil { + c.mds = inj.MDS + } + if inj.WriteLogger != nil { + c.writeLogger = inj.WriteLogger + } + if inj.Snapshotter != nil { + c.snapshotter = inj.Snapshotter + } + c.isComputeNode = inj.IsComputeNode + return nil + } +} + +type Injections struct { + MDS pilosa.MDS + WriteLogger pilosa.WriteLogger + Snapshotter pilosa.Snapshotter + IsComputeNode bool +} + // NewCommand returns a new instance of Main. func NewCommand(stdin io.Reader, stdout, stderr io.Writer, opts ...CommandOption) *Command { c := &Command{ @@ -215,12 +265,70 @@ func (m *Command) setupResourceLimits() error { return setupResourceLimitsErr } +// StartNoServe starts the pilosa server, but doesn't serve on the http handler. +func (m *Command) StartNoServe() (err error) { + // Seed random number generator + rand.Seed(time.Now().UTC().UnixNano()) + + // setupServer + err = m.setupServer() + if err != nil { + return errors.Wrap(err, "setting up server") + } + err = m.setupResourceLimits() + if err != nil { + return errors.Wrap(err, "setting resource limits") + } + + // Initialize server. + if err = m.Server.Open(); err != nil { + return errors.Wrap(err, "opening server") + } + + return nil +} + +// Register registers the node with the MDS service using whatever MDS +// implementation was injected during setup. +func (m *Command) Register() (err error) { + if m.mds == nil { + return errors.New("no MDS implementation with which to register") + } + + node := &dax.Node{ + Address: dax.Address(m.Config.Advertise), + RoleTypes: []dax.RoleType{ + dax.RoleTypeCompute, + dax.RoleTypeTranslate, + }, + } + return m.mds.RegisterNode(context.Background(), node) +} + +// CheckIn is called periodically to check in with the MDS service using +// whatever MDS implementation was injected during setup. +func (m *Command) CheckIn() (err error) { + if m.mds == nil { + return errors.New("no MDS implementation with which to check-in") + } + + node := &dax.Node{ + Address: dax.Address(m.Config.Advertise), + RoleTypes: []dax.RoleType{ + dax.RoleTypeCompute, + dax.RoleTypeTranslate, + }, + } + return m.mds.CheckInNode(context.Background(), node) +} + // Start starts the pilosa server - it returns once the server is running. func (m *Command) Start() (err error) { // Seed random number generator rand.Seed(time.Now().UTC().UnixNano()) - // SetupServer - err = m.SetupServer() + + // setupServer + err = m.setupServer() if err != nil { return errors.Wrap(err, "setting up server") } @@ -258,8 +366,8 @@ func (m *Command) UpAndDown() (err error) { // Seed random number generator rand.Seed(time.Now().UTC().UnixNano()) - // SetupServer - err = m.SetupServer() + // setupServer + err = m.setupServer() if err != nil { return errors.Wrap(err, "setting up server") } @@ -300,8 +408,8 @@ func (m *Command) Wait() error { } } -// SetupServer uses the cluster configuration to set up this server. -func (m *Command) SetupServer() error { +// setupServer uses the cluster configuration to set up this server. +func (m *Command) setupServer() error { runtime.SetBlockProfileRate(m.Config.Profile.BlockRate) runtime.SetMutexProfileFraction(m.Config.Profile.MutexFraction) @@ -368,9 +476,13 @@ func (m *Command) SetupServer() error { return errors.Wrap(err, "new stats client") } - m.ln, err = getListener(*uri, m.tlsConfig) - if err != nil { - return errors.Wrap(err, "getting listener") + if m.Config.Listener == nil { + m.ln, err = getListener(*uri, m.tlsConfig) + if err != nil { + return errors.Wrap(err, "getting listener") + } + } else { + m.ln = m.Config.Listener } // If port is 0, get auto-allocated port from listener @@ -439,6 +551,21 @@ func (m *Command) SetupServer() error { m.Config.Etcd.Dir = filepath.Join(path, pilosa.DiscoDir) } + // WriteLogger setup. + var wlw computer.WriteLogWriter = computer.NewNopWriteLogWriter() + var wlr computer.WriteLogReader = computer.NewNopWriteLogReader() + if m.writeLogger != nil { + alphaWriteLog := alpha.NewAlphaWriteLog(m.writeLogger) + wlr = alphaWriteLog + wlw = alphaWriteLog + } + + // Snapshotter setup. + var snap computer.SnapshotReadWriter = computer.NewNopSnapshotReadWriter() + if m.snapshotter != nil { + snap = alpha.NewAlphaSnapshot(m.snapshotter) + } + m.Config.Etcd.Id = m.Config.Name // TODO(twg) rethink this e := petcd.NewEtcd(m.Config.Etcd, m.logger, m.Config.Cluster.ReplicaN, version) @@ -477,6 +604,9 @@ func (m *Command) SetupServer() error { pilosa.OptServerPartitionAssigner(m.Config.Cluster.PartitionToNodeAssignment), pilosa.OptServerDisCo(e, e, e, e), pilosa.OptServerExecutionPlannerFn(executionPlannerFn), + pilosa.OptServerWriteLogReader(wlr), + pilosa.OptServerWriteLogWriter(wlw), + pilosa.OptServerSnapshotReadWriter(snap), } if m.Config.LookupDBDSN != "" { @@ -492,7 +622,6 @@ func (m *Command) SetupServer() error { } m.Server, err = pilosa.NewServer(serverOptions...) - if err != nil { return errors.Wrap(err, "new server") } @@ -500,6 +629,11 @@ func (m *Command) SetupServer() error { m.API, err = pilosa.NewAPI( pilosa.OptAPIServer(m.Server), pilosa.OptAPIImportWorkerPoolSize(m.Config.ImportWorkerPoolSize), + pilosa.OptAPIWriteLogReader(wlr), + pilosa.OptAPIWriteLogWriter(wlw), + pilosa.OptAPISnapshotter(snap), + pilosa.OptAPIDirectiveWorkerPoolSize(m.Config.DirectiveWorkerPoolSize), + pilosa.OptAPIIsComputeNode(m.isComputeNode), ) if err != nil { return errors.Wrap(err, "new api") @@ -559,7 +693,7 @@ func (m *Command) SetupServer() error { return errors.Wrap(err, "getting grpcServer") } - m.Handler, err = pilosa.NewHandler( + hndlr, err := pilosa.NewHandler( pilosa.OptHandlerAllowedOrigins(m.Config.Handler.AllowedOrigins), pilosa.OptHandlerAPI(m.API), pilosa.OptHandlerLogger(m.logger), @@ -574,7 +708,21 @@ func (m *Command) SetupServer() error { pilosa.OptHandlerRoaringSerializer(proto.RoaringSerializer), pilosa.OptHandlerSQLEnabled(m.Config.SQL.EndpointEnabled), ) - return errors.Wrap(err, "new handler") + if err != nil { + return errors.Wrap(err, "new handler") + } + + m.httpHandler = hndlr + m.Handler = hndlr + + return nil +} + +// HTTPHandler was added for the case where we want to get the full +// http.Handler, and not just those methods which satisfy the pilosa.HandlerI +// interface. +func (m *Command) HTTPHandler() http.Handler { + return m.httpHandler } // setupLogger sets up the logger based on the configuration. diff --git a/server/server_test.go b/server/server_test.go index ea4bbe524..0f2a6dfb8 100644 --- a/server/server_test.go +++ b/server/server_test.go @@ -890,6 +890,7 @@ func TestClusterExhaustingConnectionsImport(t *testing.T) { Views: map[string][]byte{ "": data, }, + SuppressLog: true, }) if err != nil { return err diff --git a/sql.go b/sql.go new file mode 100644 index 000000000..db248bf64 --- /dev/null +++ b/sql.go @@ -0,0 +1,9 @@ +package pilosa + +type SQLResponse struct { + Schema SQLSchema `json:"schema"` + Data [][]interface{} `json:"data"` + Error string `json:"error"` + Warnings []string `json:"warnings"` + ExecutionTime int64 `json:"execution-time"` +} diff --git a/sql3/planner/executionplanner.go b/sql3/planner/executionplanner.go index d1ebdd15c..1d166b4d4 100644 --- a/sql3/planner/executionplanner.go +++ b/sql3/planner/executionplanner.go @@ -4,7 +4,6 @@ package planner import ( "context" - "encoding/json" pilosa "github.com/featurebasedb/featurebase/v3" "github.com/featurebasedb/featurebase/v3/batch" @@ -86,18 +85,6 @@ func (p *ExecutionPlanner) CompilePlan(ctx context.Context, stmt parser.Statemen rootOperator, err = p.optimizePlan(ctx, rootOperator) } - // Log the plan. This happens even if an error occurred. - switch rootOperator.(type) { - case *PlanOpInsert: - // Don't log the insert plan since it can be very large. - case nil: - // pass - default: - plan := rootOperator.Plan() - a, _ := json.MarshalIndent(plan, "", " ") - p.logger.Debugf(string(a)) - } - return rootOperator, err } diff --git a/sql3/planner/executionplannersystemtables.go b/sql3/planner/executionplannersystemtables.go index b0fe25770..85623c91c 100644 --- a/sql3/planner/executionplannersystemtables.go +++ b/sql3/planner/executionplannersystemtables.go @@ -11,6 +11,9 @@ import ( "github.com/pkg/errors" ) +// Ensure type implements interface. +var _ pilosa.SchemaAPI = (*systemTableDefintionsWrapper)(nil) + type systemTableDefintionsWrapper struct { schemaAPI pilosa.SchemaAPI } @@ -53,6 +56,10 @@ func (s *systemTableDefintionsWrapper) IndexInfo(ctx context.Context, indexName return i, nil } +func (s *systemTableDefintionsWrapper) FieldInfo(ctx context.Context, indexName, fieldName string) (*pilosa.FieldInfo, error) { + return nil, pilosa.ErrNotImplemented +} + func (s *systemTableDefintionsWrapper) Schema(ctx context.Context, withViews bool) ([]*pilosa.IndexInfo, error) { schema, err := s.schemaAPI.Schema(ctx, withViews) if err != nil { diff --git a/sql3/test/defs/defs_keyed.go b/sql3/test/defs/defs_keyed.go index dd3ed858e..30d031e54 100644 --- a/sql3/test/defs/defs_keyed.go +++ b/sql3/test/defs/defs_keyed.go @@ -13,8 +13,8 @@ var keyed = TableTest{ ), srcRows( srcRow("one", int64(11), []int64{11, 12, 13}, int64(101), "str1", []string{"a1", "b1", "c1"}), - srcRow("two", int64(22), []int64{21, 22, 23}, int64(201), "str2", []string{"a2", "b2", "c2"}), - srcRow("three", int64(33), []int64{31, 32, 33}, int64(301), "str3", []string{"a3", "b3", "c3"}), + srcRow("two", int64(22), []int64{11, 12, 23}, int64(201), "str2", []string{"a2", "b2", "c2"}), + srcRow("three", int64(33), []int64{11, 32, 33}, int64(301), "str3", []string{"a3", "b3", "c3"}), srcRow("four", int64(44), []int64{41, 42, 43}, int64(401), "str4", []string{"a4", "b4", "c4"}), ), ), @@ -36,8 +36,8 @@ var keyed = TableTest{ ), ExpRows: rows( row("one", int64(11), []int64{11, 12, 13}, int64(101), "str1", []string{"a1", "b1", "c1"}), - row("two", int64(22), []int64{21, 22, 23}, int64(201), "str2", []string{"a2", "b2", "c2"}), - row("three", int64(33), []int64{31, 32, 33}, int64(301), "str3", []string{"a3", "b3", "c3"}), + row("two", int64(22), []int64{11, 12, 23}, int64(201), "str2", []string{"a2", "b2", "c2"}), + row("three", int64(33), []int64{11, 32, 33}, int64(301), "str3", []string{"a3", "b3", "c3"}), row("four", int64(44), []int64{41, 42, 43}, int64(401), "str4", []string{"a4", "b4", "c4"}), ), Compare: CompareExactUnordered, @@ -59,8 +59,8 @@ var keyed = TableTest{ ), ExpRows: rows( row("one", int64(11), []int64{11, 12, 13}, int64(101), "str1", []string{"a1", "b1", "c1"}), - row("two", int64(22), []int64{21, 22, 23}, int64(201), "str2", []string{"a2", "b2", "c2"}), - row("three", int64(33), []int64{31, 32, 33}, int64(301), "str3", []string{"a3", "b3", "c3"}), + row("two", int64(22), []int64{11, 12, 23}, int64(201), "str2", []string{"a2", "b2", "c2"}), + row("three", int64(33), []int64{11, 32, 33}, int64(301), "str3", []string{"a3", "b3", "c3"}), row("four", int64(44), []int64{41, 42, 43}, int64(401), "str4", []string{"a4", "b4", "c4"}), ), Compare: CompareIncludedIn, @@ -84,10 +84,127 @@ var keyed = TableTest{ hdr("a_string_set", fldTypeStringSet), ), ExpRows: rows( - row("two", int64(22), []int64{21, 22, 23}, int64(201), "str2", []string{"a2", "b2", "c2"}), + row("two", int64(22), []int64{11, 12, 23}, int64(201), "str2", []string{"a2", "b2", "c2"}), ), Compare: CompareExactUnordered, SortStringKeys: true, }, }, + PQLTests: []PQLTest{ + { + name: "minrow", + Table: "keyed", + PQLs: []string{"MinRow(field=an_id_set)"}, + ExpHdrs: hdrs( + hdr("an_id_set", fldTypeID), + hdr("count", fldTypeID), + ), + ExpRows: rows( + row(int64(11), int64(1)), + ), + }, + { + name: "maxrow", + Table: "keyed", + PQLs: []string{"MaxRow(field=an_id_set)"}, + ExpHdrs: hdrs( + hdr("an_id_set", fldTypeID), + hdr("count", fldTypeID), + ), + ExpRows: rows( + row(int64(43), int64(1)), + ), + }, + { + name: "topk", + Table: "keyed", + PQLs: []string{"TopK(an_id_set, k=2)"}, + ExpHdrs: hdrs( + hdr("an_id_set", fldTypeID), + hdr("count", fldTypeID), + ), + ExpRows: rows( + row(int64(11), int64(3)), + row(int64(12), int64(2)), + ), + }, + { + name: "topn", + Table: "keyed", + PQLs: []string{"TopN(an_id_set, n=2)"}, + ExpHdrs: hdrs( + hdr("an_id_set", fldTypeID), + hdr("count", fldTypeID), + ), + ExpRows: rows( + row(int64(11), int64(3)), + row(int64(12), int64(2)), + ), + }, + { + name: "rows", + Table: "keyed", + PQLs: []string{"Rows(field=an_id_set)"}, + ExpHdrs: hdrs( + hdr("an_id_set", fldTypeID), + ), + ExpRows: rows( + row(int64(11)), + row(int64(12)), + row(int64(13)), + row(int64(23)), + row(int64(32)), + row(int64(33)), + row(int64(41)), + row(int64(42)), + row(int64(43)), + ), + }, + { + name: "includescolumn", + Table: "keyed", + PQLs: []string{"IncludesColumn(Row(an_id_set=12), column='two')"}, + ExpHdrs: hdrs( + hdr("result", fldTypeBool), + ), + ExpRows: rows( + row(true), + ), + }, + { + name: "constrow", + Table: "keyed", + PQLs: []string{"Extract(ConstRow(columns=['two']), Rows(an_id))"}, + ExpHdrs: hdrs( + hdr("_id", fldTypeString), + hdr("an_id", fldTypeID), + ), + ExpRows: rows( + row("two", int64(201)), + ), + }, + { + name: "fieldvalue", + Table: "keyed", + PQLs: []string{"FieldValue(field=an_int, column='three')"}, + ExpHdrs: hdrs( + hdr("value", fldTypeInt), + hdr("count", fldTypeInt), + ), + ExpRows: rows( + row(int64(33), int64(1)), + ), + }, + { + name: "unionrows", + Table: "unkeyed", + PQLs: []string{"Count(UnionRows(Rows(field=an_id_set)))"}, + ExpHdrs: hdrs( + hdr("count", fldTypeID), + ), + ExpRows: rows( + row(int64(4)), + ), + }, + }, } diff --git a/sql3/test/defs/defs_unkeyed.go b/sql3/test/defs/defs_unkeyed.go index daf7254d2..038870480 100644 --- a/sql3/test/defs/defs_unkeyed.go +++ b/sql3/test/defs/defs_unkeyed.go @@ -94,4 +94,17 @@ var unkeyed = TableTest{ SortStringKeys: true, }, }, + PQLTests: []PQLTest{ + { + name: "options", + Table: "unkeyed", + PQLs: []string{"Options(Count(Row(an_id_set=1)), shards=[0])"}, + ExpHdrs: hdrs( + hdr("count", fldTypeID), + ), + ExpRows: rows( + row(int64(0)), + ), + }, + }, } diff --git a/sql3/test/defs/types.go b/sql3/test/defs/types.go index fbb93962e..10271e765 100644 --- a/sql3/test/defs/types.go +++ b/sql3/test/defs/types.go @@ -39,6 +39,7 @@ type TableTest struct { name string Table source SQLTests []SQLTest + PQLTests []PQLTest } // Name returns a string name which can be used to distingish test runs. It @@ -102,6 +103,26 @@ func (s SQLTest) Name(i int) string { return name } +type PQLTest struct { + name string + PQLs []string + Table string + ExpHdrs []*planner_types.PlannerColumn + ExpRows [][]interface{} + ExpErr string +} + +// Name returns a string name which can be used to distingish test runs. It +// takes an integer value which will be used as part of a generic table name in +// the case where a name value was not provided in the definition. +func (s PQLTest) Name(i int) string { + name := fmt.Sprintf("test-%d", i) + if s.name != "" { + name = s.name + } + return name +} + // The following "source" types are helpers for creating a test table. type sourceColumn struct { name string diff --git a/test/pilosa.go b/test/pilosa.go index bb105efe1..a3d45eb44 100644 --- a/test/pilosa.go +++ b/test/pilosa.go @@ -295,6 +295,36 @@ func Do(t testing.TB, method, urlStr string, body string) *httpResponse { return &httpResponse{Response: resp, Body: string(buf)} } +// Do executes http.Do() with an http.NewRequest(). +func DoProto(t testing.TB, method, urlStr string, body []byte) *gohttp.Response { + t.Helper() + req, err := gohttp.NewRequest( + method, + urlStr, + bytes.NewReader(body), + ) + if err != nil { + t.Fatal(err) + } + + req.Header.Set("Content-Type", "application/x-protobuf") + req.Header.Set("Accept", "application/x-protobuf") + + // set a timeout instead of allowing gohttp.Defaultclient to + // potentially hang forever. + hc := &gohttp.Client{ + Timeout: time.Second * 30, + } + resp, err := hc.Do(req) + + if err != nil { + fmt.Printf(" hc.Do() err = '%v'\n", err) + t.Fatal(err) + } + + return resp +} + func CheckGroupBy(t *testing.T, expected, results []pilosa.GroupCount) { t.Helper() if len(results) != len(expected) { diff --git a/test/transaction.go b/test/transaction.go index 2f91ff7d1..7f7b7f957 100644 --- a/test/transaction.go +++ b/test/transaction.go @@ -6,7 +6,7 @@ import ( "testing" "time" - pilosa "github.com/featurebasedb/featurebase/v3" + pilosa "github.com/molecula/featurebase/v3" ) const deadlineSkew = time.Second diff --git a/transaction_test.go b/transaction_test.go index 73024b2fc..d9f071597 100644 --- a/transaction_test.go +++ b/transaction_test.go @@ -8,9 +8,9 @@ import ( "testing" "time" - pilosa "github.com/featurebasedb/featurebase/v3" - "github.com/featurebasedb/featurebase/v3/logger" - "github.com/featurebasedb/featurebase/v3/test" + pilosa "github.com/molecula/featurebase/v3" + "github.com/molecula/featurebase/v3/logger" + "github.com/molecula/featurebase/v3/test" ) // TestTransactionManager currently uses an in memory transaction diff --git a/txfactory.go b/txfactory.go index 00ee963c5..c801c2ed3 100644 --- a/txfactory.go +++ b/txfactory.go @@ -175,13 +175,15 @@ func (f *TxFactory) NewQcx() (qcx *Qcx) { Grp: f.NewTxGroup(), Txf: f, } - if f.holder != nil && f.holder.executor != nil { - qcx.workers = f.holder.executor.workers - } if f.typeOfTx == "roaring" { qcx.isRoaring = true } - _ = testhook.Opened(f.holder.Auditor, qcx, nil) + if f.holder != nil { + if f.holder.executor != nil { + qcx.workers = f.holder.executor.workers + } + _ = testhook.Opened(f.holder.Auditor, qcx, nil) + } return }