From 04aa8b18fc49107d23532259f0ee4f83eeb6f30b Mon Sep 17 00:00:00 2001 From: Travis Date: Wed, 9 Mar 2022 09:15:55 -0600 Subject: [PATCH] Add DAX - full list of squashed commits below In this commit, the Directive is mocked; it doesn't actually reach out to a controller. Limits key translation to only those partitions (per index) specified in the Directive. Attempting to create or find a key (or ID) for a partition which is not handled by this node will result in an error; translation requests are no longer forwarded to other nodes. Limits import into only those shards specified, per index, by the Directive. Attempting to import into a shard which is not handled by this node will result in an error. Stub out /directive endpoint The `applyDirective()` method still needs to be implemented. Update mds references to use the new /mds/types structure In mds, we moved the shared types to mds/types. FeatureBase needs to reference those instead. This also bumps the mds version in go.mod. Implement the Add/Remove Index part of Holder.ApplyDirective() This adds functionality to `Holder.ApplyDirective()` which adds or removes indexes (tables) based on those provided in the Directive. Still to be implemented here: shards and partitions. WIP: remove client from Batch Move Batch into its own package: batch Also, in order to avoid import loops, this introduces packages: /batch/types /client/types Reorganize the Importer-related code Moved the Importer interface to package: batch Move the "pilosa client" implementation of the Importer interface to package: client Modify batch.NewBatch to take an Importer (not client) This commit modifies the batch.NewBatch() function to use a functinal option on Batch to inject an Importer into the Batch. Prior to this, NewBatch() took a pointer to a client, which was a little too restrictive. Now, MDS can implement an Importer which uses information from MDS to determine to which node(s) the import calls should be directed. Add client.SetAuthToken() method to satisfy SchemaManager interface Update ApplyDirective logic to include fields. This needs more work, but it was enough to get a basic test passing. Move Transaction type into /types package. Add interface check on batch.Importer no-op implementation Updated ApplyDirective to create all currently support Field types There are still the following TODOs: - [ ] impolement field options (ex: decimal scale, int min/max, etc). - [ ] `time` fields Added support for Decimal.Scale in ApplyDirective Update mds dependency Add /health endpoint Update to use dax (dax/mds) instead of mds. After moving the mds repository into the dax repository as a sub-package, this commit changes everything in FeatureBase to use the dax repo instead of the now abandoned mds repo. Introduce and implment the WriteLogger interfaces. This adds both a `WriteLogReader` and `WriteLogWriter` interface. They are both implemented by the implementation: `fileWriteLogger`. The `fileWriteLogger` uses the dax/writelogger API to append log messages to files on disk. Add WriteLogWriter.ImportRoaring method to interface This commit adds the `ImportRoaring` method to the `WriteLogWriter` interface. Still to implement are the `Import` and `ImportValue` methods. Reorganize the ApplyDirective code The primary goal was to cache the incoming Directive on the Holder prior to applying all of the changes in the directive (i.e. loading data from the WriteLogger) because applying those changes often validated against the accepted state of the node. Implement all of the WriteLogger read/write methods Implement the HTTP WriteLogger implementation WIP: Introduce shard.Version. Implement snapshotter. Add HTTP Snapshotter implementation This also recofigures server to use the HTTPSnapshotter instead of the FileSnapshotter. Implement snapshotter: TableKeys Implement snapshotter: FieldKeys Dependency dance last of the dependency dance Add support for prototype This adds the Makefile targets to build the docker container and push it to ECR. SQL3 changes which break with dax changes Missed TODO: implement FieldVersion version to WriteLogger Address bug causing missing TranslateStores to error Originally, we tried to limit the TranslateStores which get allocated to only those for which the node is responsible. This works when adding a new table. But if a table already exists, there's no logic to start missing TranslateStores. This reverts back to the old FeatureBase logic which brutishly allocates a TranslateStore for every partition, even if one is not needed. We need to address this by allowing the ApplyDirective logic to initialize TranslateStores when they don't yet exist. Move the ImportRoaringShardRequest type to the types package Since the ImportRoaringShardRequest object is part of the Importer interface, we need to move it to a non-root (i.e. pilosa) package. All the other interface types are either concrete types or part of a sub-package (such as roaring). We do this to prevent an implementer of the interface from having to import the entire pilosa package and risk circular imports. buncha changes to support latest dax stuff Move dax related types to /dax sub-package This commit moves all the common "dax" types into the /dax sub-package. The idea is to ensure that featurebase does not import dax at all. It's ok if dax imports featurebase. In the future, we might need to split the dax sub-package (common data types used by muliple molecula data-plan services) into it's own repo. Add type: dax.Schema This isn't currently being used; I started to use is as a replacement for pilosa_client.Schema, but then deferred that. But we'll need to do it eventually, so it doesn't hurt to have this type in place. Export RowIDs.Merge() method for use in orchestrator. Add CreateSQL method to dax.Table type The CreateSQL() method will return the "CREATE TABLE" statement required to create the dax.Table. Comment out confusing writelogger log message. We need to revisit this, but for now, this log message is confusing. Also, rename daxSharder to versionStore. Remove hard-coded AWS account Implement more FieldOptions such as Epoch Some of the FieldOption logic was stubbed out in the dax package. This commit fills that out more; specifically, it adds the dax.Field.Options.Epoch parameter. export stuff needed for TopK in orchestrator export ValCount stuff to implement Percentile in orchestrator export more stuff to support less code in orchestrator, shared objs Port dax repo over to featurebase/dax (run all as sub-services) This commit does ALOT. Sorry. It introduces a `featurebase dax` sub-command which can be configured to run the various dax services as sub-services within the same process, or individually as the lone service in process. It also changes all the URL paths to be prefixed with the service name. So for example, instead of calling localhost:8080/status, you would now call localhost:8080/featurebase/status. Also, note that all services provide a /health endpoint to confirm they are running in process. Clean up integration tests. Remove PILOSA_ config prefix. Remove duplicate clients (mistake from porting dax to featurebase) Rename sub-service "featurebase" to "computer" In the places where we have hard-coded the sub-service name into a URI path, I've tried to tag the line with a comment containing: `// #SERVICEPATHPREFIX` Update copilot manifest files to reference "computer" Port dax/README.md from dax repository Separate (toml) Queryer Config from Injections We needed to separate the toml config from the configuration required to inject sub-services into the Queryer. I'm not sure this is the best solution, but it's *a* solution. So here we are. Clean up (i.e. remove) the queryer "implementations" package Remove old test file Run WriteLogger and Snapshotter as local sub-services. Prior to this commit, the writelogger and snapshotter services only worked when run as separate services. This allows them to be run in the same process as all the other dax services. There is still some naming issues that we should address, but it's functional for now. Clean up (i.e. organize) the intra-service interfaces. Implement alpha Director for local messages from MDS to Computer Prior to this commit, messages from MDS to the computer service were still going over http. This commit introduces an interface implementation which registers the local computer command, and use that command's API to directly reference methods used by the Director. Clean up a few more interface names Add Queryer OpenAPI document. Update copilot manifests to reflect latest changes Add OpenAPI documents for WriteLogger and Snapshotter Add OpenAPI document for MDS service Add OpenAPI document for Computer service Consolidate errors to use fb/errors package. This commit is a first pass at trying to ensure that all of the DAX code uses: "github.com/molecula/featurebase/v3/errors" This package is a wrapper for "github.com/pkg/errors", so going forward we want to avoid importing that package. The only method which isn't backward-compatible is `New()`; the New() method in the featurebase/errors package takes an errors.Code. If this becomes a problem, we could change this by reverting New() and then introducing something like NewCoded(). But for now I think it might actually discourage someone from just creating a New() error without thinking about how it should be coded. Introduce VersionStore interface Move the existing VersionStore code to the `inmem` package as the in-memory implementation of the new dax.VersionStore interface. Introduce NodeService interface With this, the Controller can maintain a registry of nodes by using this NodeService interface as opposed to an in-memory map of nodes on the Controller struct. This also adds an inmem implementation of the NodeService interface. Introduce controller.Balancer interface This moves the existing balancer package to controller/naive package. The idea is to allow us to add a different Balancer implementation in the future. Introduce DirectiveVersion interface This commit also includes *A LOT* of refactoring to use dax.Worker and dax.Job types everywhere instead of strings. Introduce Schemar interface The previous `Schemar` struct was moved to the `schemar/inmem` package, and `Schemar` is now an interface implemented by that inmem package. Remove unused type `nUnit` Add boltdb implementation of VersionStore interface. This removed the previous sqlite implementation; we decided not to use sqlite for now (as a basic, local disk implementation) because it requires CGO. -------------------------------------------- No longer applicable: Add sqlite implementation of VersionStore interface. This commit implements the VersionStore interface using sqlite. Sqlite requires CGO, so this may not be something we want to include, but it's implemented here to get a feel for how an external implementation might be used; the next step will be to determine how the user configured FeatureBase to run using sqlite as a backing store for services like MDS. Add boltdb implementation of NodeService and DirectiveVersion interfaces. Add boltdb implementation of naive Balancer interfaces. This includes the two interfaces defined in `naive/balancer.go`: - WorkerJobService - FreeJobService Add boltdb implementation of Schemar interface. clean up a linter issue Thread context.Context through all the interfaces. Some of the interface implementations are going to use context, so we need to make that part of the interface. The boltdb implementations, for example, take a context. This is probably so we can do things like cancel or timeout operations. Update interfaces to return error; remove `panic(err)` everywhere. Down-rev grpc version to 1.38.0 Later versions (after 1.42.0?) cause MustRunCluster.Close() in tests to deadlock. This commit also adds an `isComputeNode` feature flag around some of the write log and shard/partition check functionality so that it doesn't run under normal conditions (this is excercised by running the sql3 tests for example). Add MDS_Persistence test to cover meta data persistence This adds a basic test which configures the MDS container to use boltdb as its persistence storage, saved on a docker volume. Then, the mds container is stopped/replaced, and we confirm that the data stored on the volume is availble to the new MDS container. Fix a few things after rebase with sql-experiment branch The lastest version of sql-experiment contains a fairly significan refactor of the way query iteration works. This commit adjusts for those changes. pull dax IDK changes in to FB IDK (#2177) * pull dax IDK changes in to FB IDK * Move docker-related IDK build stuff to featurebase root Building the docker image required the root level go.mod and vendor directory. This change moves the make targets to the root level Makefile, and the Dockerfiles now copy the root level vendor directory (and everything else in the root for that matter). * Fix batch- and client-related tests * InitializePoller on MDS restart/replacement Prior to this change, if MDS was restarted, its internal poller (which maintains an in-memory list of nodes to poll) is empty. This is bad, because it doesn't know about nodes that it should be polling. This change fixes that. Upon MDS startup, it intializes the poller with the list of nodes that MDS keeps in persistent storage (currently: boltdb). * Add EFS volume to MDS Copilot manifest This allows us to use MDS's persistent storage (via boltdb) in the Copilot demo by saving metadata in a boltdb file on EFS. * Thread logger.Logger through all dax components * Revert some of the breaking changes from DAX development. When we first started prototyping DAX, we made changes to the featurebase core code which would have broken the existing featurebase functionality. This commit reverts some of those changes. Anywhere that we need to modify core featurebase functionilty, we put it behind some kind of feature flag. This flag is typically determined by whether the running node is a "compute" node (i.e. DAX.COMPUTER.RUN = true). Co-authored-by: Travis Turner add packaging for DAX need cgo for datagen build bind to 0.0.0.0, pass GOOS and GOARCH explicitly not sure if the explicit GOOS/GOARCH is actually necessary... Get INSERT INTO (aka ingest) working through SQL3 This commit does a few things which I'll try do describe here. - Introduces a Qctx interface. The existing Qcx is an implementation of this interface, and can be used exactly how it has been. But this allows us to abstract away the notion of Qcx in the Queryer (which is handling SQL3) until we're ready to address that. As an example, the Qcx has a notion of a featurebase Holder, but that doesn't make sense when we're at the Queryer layer. For now, the Qctx used in the Queryer is a no-op. - Adds a ComputeAPI interface implementation for the Queryer. This is effectively the Import() and ImportValues() methods used for ingest. The logic here handles the incoming ImportRequest by first doing any necessary column and row translation for the entire request, then it splits the records by shard, and generates a new ImportRequest per shard with only the shard-appropriate records. - Changes the mds.Importer to take an MDS interface implementation (which can be an mds client) instead of an mdsAddress. This allows us to use a localy MDS implementation rather than assuming we need a client to make calls over a network. Add queryer.Importer interface to handle ingest via SQL (#2203) * Add queryer.Importer interface to handle ingest via SQL This is meant to support ingest through SQL when the queryer and the compute services are running in the same process, or when they are on seperate processes and need to talk via http client. * remove datagen from RPM was originally added as a convenience to generate test data, but is unused and annoying because datagen doesn't easily cross-compile due to cgo * add marshalUnmarshal to controller to avoid passing pointers passing pointers across API boundaries can cause unpredictable things in local vs remote configurations. Co-authored-by: Matthew Jaffee "fix" a few issues with wrong default partition numbers these still need to be properly fixed and actually get the correct data from MDS go mod tidy Introduce TableQualifier (OrganizationID/DatabaseID) (#2220) * add check in ApplyDirective that version is increasing fix TestAPIDirective to make version always increasing * fix docker image build and break out dax test in CI We have to run the DAX integration tests separately as they call out to Docker, and so it isn't easy to run them in a Docker container as the other tests do. So we run them directly on the CI runner which has Docker and Go installed. We also explicitly exclude these tests from running during the other tests. Also my editor was automatically reformatting some comments badly which is why I added the "data" thing in those two places * add timeout to poller * give Poller a default Logger apparently we can NPE sometimes, seen in CI: https://gitlab.com/molecula/featurebase/-/jobs/3028286364 * bunch of testing fixes, mostly IDK/DAX related make MDS error if sendDirectives errors, don't just log. sendDirectives can error if computer nodes disagree about the validity of a schema (for example), in which case it might need to get deleted and user notified somehow. very messy, needs more thought. re-introduce old env prefix to maintain compatibility with master branch make self-contained dax container for IDK testing build IDK images from source (now that all the source is available since it's in the same repo) catch errors in DoExtractQuery in idktest.go fix IDK bug where prefix path was hardcoded in all cases rather than only when useMDS was true fix TestBatchTargetMDS... needed to add field options and catch error when creating table. also needed an _id field * fix env prefix in tests * WIP getting tests to pass, wanna see CI * don't error if we get a zero version directive and we don't have a directive yet * cleanup debugging junk * "fix" future.rename thing, run IDK tests * Introduce TableQualifier (OrganizationID/DatabaseID) This commit introduces a lot of new types (in dax/table.go) related to TableQualifer (which is made up of OrganizationID and DatabaseID), as well as things like TableID and TableKey. For the most part, we try to thread a QualifiedTableID through the entirety of DAX. There are some places (for example in the Balancers, which are just aware of string keys) which use a string TableKey (tbl__org__db__tableid). * Remove some debugging comments * Add Org/DB support to CLI. This commit adds support for special commands: SET SET ORG acme SET DB db1 USE db1 * remove ".pulled" from IDK Makefile I don't think we need it any more as most things can be built locally. I think it was only there to refresh the FeatureBase images that were tagged as master, but we don't need to do that any more. * Change DAX json tags to kebab-case (i.e. hyphenated) This commit also renames some struct arguments to more accurately reflect their type: for example, renaming `Table` to `TableKey` when the type is TableKey. * Return DAX TableName in SHOW TABLES (instead of Index.Name) There are cases where SchemaAPI is used to return DAX friendly table names (as opposed to featurebase index names, which are DAX TableKey). This is an attempt to do that. With that said, it's not ideal because anything could call those API methods and expect the other type. * Fix a bug which wasn't completely dropping a table. When using boltdb as a backend, DROP TABLE wasn't removing the reverse-lookup key for the table in boltdb. * Remove idk/testenv/certs which got accidentally committed. also update .gitignore to include those. * Fix IDK ingest tests to be TableQualifier aware. * Add example Table types to dax/table.com godoc. * ignore idk.Main fields for flags, upgrade commandeer * go mod tidy * Fix DAX integration tests: ingester using wrong ENV VARs We change from ORGANIZATION_ID to ORG_ID and from DATABASE_ID to DB_ID * Clarify things around idk (docker) tests * Stop running TestKafkaSourceIntegration with t.Parallel() This test can't be run in parallel as it's currently written. Doing so allows for interleaving of messages to the same kafka topic between tests. I didn't attempt to modify the test so it could be run in parallel. That could be done, but left for someone more ambitious. Co-authored-by: Matthew Jaffee Require Directive.Version be a non-zero value. (#2227) Because the directive cached on the holder is not a pointer, its default version is 0. In order to avoid having to compare against that, we just require that Directive.Version start at 1. General, non-invasive code cleanup and comment adjustment. Move ImportRoaringShardRequest out of the types package Early on in the DAX development, I moved ImportRoaringShardRequest into a types package. There must have been some import loop going on, but since that is not longer the case, it's safe to move this back into the core featurebase (er... pilosa) package. Move Transaction struct back into the pilosa package (from types) Revert some name changes (cli -> client) Add DAX Handler CloseTimeout This was implemented in htt_handler.go, but it had been commented out in the DAX handler. This just uncomments that and finishes the implementation. Remove Qcx from queryer.Importer interface This sets us up to revert the Qctx interface that was initially introduced to allow us to abstract away the need for a Qcx when calling the ComputeAPI from a remote service (i.e. the queryer). Add some go-doc comments and remove unused code. Move SchemaManager setup from datagen to idk.Main (#2233) The set for idk.SchemaManager (for dax implementations) was previously in datagen. This may have been because of some import loop problem during development, but that's no longer an issue. The setup for this should be in idk.Main so anything using that can leverage the MDS-specific SchemaManager setup. Fix issues around nil TxFactory First, don't return a nil. Rather return a new *TxFactory (with no holder). Second, don't call `f.holder` in the testhook outside of checking if `f.holder` is nil. Wrap all bare errors Make service prefixes constants Instead of having `"computer"` throughout the code, use instead a constant: `dax.ServicePrefixComputer`. MDS skip errors when sending empty directives also add in the docker-login and ecr-push changes for serverless DAX Fix the logic in Directive.IsEmpty() (#2236) Update the cached value for Index.translatePartitions In the case where a node already knows about an index, but its assignment of partitions for that index changes (for example, when another node goes down and the node in question is now responsible for more partitions than it previously was), then we need to update the cached value of Index.translatePartitions because that's used in translation checks. minor fixes for IDK-related bugs WIP: tokenize CLI to access cloud FB CLI cloud support with automatic token refresh Also adds support for a GET command which allows making HTTP GET queries to cloud CP which can be handy for debugging stuff. E.g. GET /v2/databases buncha little fixes working on writelogger stuff fix writelogger/snapshotter setup bugs implement writelogging for importRoaringShard add debug endpoint to MDS use shard transactional endpoint in MDS datagen add debugging to API related to writelogger revert handleroption change clean up big PR remove "GET" command from CLI for making arbitrary HTTP request to cloud control plane (was a messy hack and not that useful) remove json tags from FB objects where we had to duplicate the object elsewhere due to import loops and weren't actually json encoding it unexport handlerOption which was exported to try to avoid doing certain things if we're in DAX mode, but I didn't end up merging that code. remove (hopefully) unecessary extra call to api.indexField fix some formatting, unexport some vars, godoc, etc oops, fix build failure Update FeatureBase CLI to support a standard deployment The standard deployment uses a different endpoint and request payload. This commit tries to detect is the standard deployment is being used, and if so, it uses a standard-specific FBQueryer. It also modifies the auto-detection logic to try standard featurebase and dax ports in the case where a port was not provided. MDS API refactor (#2259) * MDS API refactor table IDs are exposed but only created server side also cleaned up dax Makefile * clean up review feedback Co-authored-by: Travis Turner * remove TablesByName * rip out inmem implementations and use boltdb everywhere * remove inmem balancer, create bolt tempfile by default on startup * WIP on snapshot table impl and test * Minor comment and code layout adjustments. This commit also adds the `Equals` method to `QualifiedTableID` for equality comparisons. It's no longer safe to compare struct (two structs might still be equal even if one of the structs doesn't have a `Name` value. * Use a unique docker network for each dax test Ocassionally we would see some test failures due to a network already existing. This shouldn't happen, but to avoid that, this commit generates a unique name for each sub test (which gets deleted at the end of every test). * Fix one instance of NewQualifiedTableID losing Name We should probably check the other instances and see if Name is getting lost. * simplify unique network stuff and fix api directive tests * Remove TableIDRequest and TableIDResponse types for /table-id (#2267) For the mds/table-id http requests, just use dax.QualifiedTableID as both the request and response types. * remove lattice from dax, no error on node re-reg, dax docker-compose * various updates * WIP: mds-refactor branch review * no-op on SnapshotTableKeys if table is not keyed * Makefile helpers * add doWeCare so controller doesn't fail unnecessarily * clean up table creation (#2272) * Strip underscores from TableID stub name * fix boltdb versionstore tests: generate unique, sorted tables * fix controller test related to reregistering a node * JobSet -> generic Set Co-authored-by: Travis Turner Co-authored-by: Travis Turner Cleanup after rebase on master The latest rebase on master entailed all the client/batch changes as well as some of the qcx refactoring. It made for a hairy rebase. This commit fixes some of the tests that were failing after that rebase. Fix batch/client import loop missed during rebase (#2280) It's not surprising that `batch` can't import `client`. It was doing that here (importing an error type from the `client` package). What is surprising is that it's okay for `batch_test.go` to import `client` even though `batch_test.go` is an internal test and therefore part of the `batch` package. different boltDB's for schemar/controller, explicit balancers nice helpers for dax docker-compose, make build really fast build FB binary outside of docker, then create Docker image with its working dir in an empty subdirectory so it doesn't send a GB of context to the daemon. error on unassigned jobs and use client with timeout fix CR feedback deregister batch of nodes also make removal faster via director dial timeout implement WorkersForJobPrefix so orchestrator doesn't make up shards also fix some godocs and remove unused method Run sub-tasks of a Directive concurrently in a worker pool. (#2275) * Run sub-tasks of a Directive concurrently in a worker pool. This allows the compute node to concurrently load shapshot and writelog data concurrently, instead of one keyset/partition/shard at a time. It introduces a config parameter called `DirectiveWorkerPoolSize`. * code review cleanup * Use unique container names in DAX integration tests We were seeing "container already exists" errors in CI, so just to be safe, this commit constructs a unique container name for every container in the DAX integration test run. Stub in SystemAPI to Queryer (note: will not work if used) This just makes is so that dax can compile. Actually implementing system-table functionality for dax will take some planning. Tlt/dax merge prep (#2282) * Remove copilot directory * Remove Dockerfile-datagen-long * Remove orphaned RegisterNodeRequest This type is not defined in the dax/mds/http package. * implement TIMEQUANTUM and TTL in Table.Field type * Remove the "service" misdirection in queryer/writelogger/snapshotter. We had originally used an additional layer, er.. package, for a "service". The main distinction was that the Config differed in that it was internal, unlike the Config that we need to provide for the top-level server config (i.e. toml). Having that additional layer just to support a different Config seemed premature at best. So I'm removing it. * Remove dax docker containers no longer used in tests Since we run everything as "featurebase", we don't have multiple container types anymore. * Some minor comment updates * Remove nfpm stuff related to dax * Fix linter issues Fix "duplicate" issues raised by sonarcloud. run docker components of dax integration tests with coverage trying to get dax integration coverage add coverate volume mounts throughout dax integration tests add a lock, tweak dax Makefile, remote flag on query handler remove some unused code convert batch tests to use clustertests to get coverage maybe fix clustertests more authclustertests fixes, test is failing locally but also seems to have been silently failing in CI prior to these changes... let's see if it's still silent fix some lint to kick CI just re-running the job wasn't working... strange behavior remove RetryLogic test and pipe which don't work RetryLogic test removed due to etcd changes. Seebs thinks we shouldn't test this here. Pipe was being ignored since we're no longer using "bash -c" to execute the command. If we need to generate that output file we'll either have to reintroduce bash -c and set -o pipefail so that it actually fails properly, or figure out some other solution. shooting into the dark... first cut at bulk node registration remove unused stuff from batch tests, set coverpkg to ../... batch registration timeout and fix tests disable most tests and don't run fb background batch test debuggin!!!!!!!!! and then he tried this.... Implement importer (for INSERT INTO) in the Queryer Prior to this, we we passing a nil value in for the importer to the planner.NewExecutionPlanner in the Queryer. This meant that INSERT INTO statements didn't work. Now they should. It uses the importer that we build for IDK in /idk/mds/importer.go, and wrapps that with a type that can determine if the provided string "index" is of the form indexName or TableKey. turn off debug mode, fix log saving Run sql3 test definitions in a dax integration test There are currently 22 tests which are not passing. They are skipped in the "skips" slice. WIP, not working, pql queries to tests Add TableQualifier to PQL query logic in the Queryer Add more PQL tests to the keyed table Allow instant node registration if registration-batch-timeout=0 When running dax services in process, we don't want to wait 3s for the compute node to register; we know it's there because it's in the same process. Fixes related to IncludesColumn PQL test. Tests for ConstRow and FieldValue cleanup add UnionRows and Options, better error reporting on bad queries delete unused schemar client.go, clean up unused in batch test CI move test timeouts into more reasonable territory apparently this had already been done, but got merge-stommped at some point move dax bolt test helpers into dax package Add computer CheckIn routine (#2296) * Add computer CheckIn routine This adds a background routine which sends a "check-in" request to MDS every . This is to address the case where the poller has removed a computer node from the node list (due to a network fault, for example), but the node is still healthy and becomes available again. In that case, the node needs to "check-in" to tell MDS it is still there. MDS will likely send the node a new directive with Method=reset telling the node to delete all of its data an apply the latest directive. * Don't send directives to Deregistered (i.e. removed) nodes We have an issue where we're locking on sendDirective in the controller, and when the node is unavailable, the send hangs and never releases the lock. This is a temporary fix for that until we address the real problem. Fix .gitlab-ci.yml after rebase fix some indentation shenanigans (cherry picked from commit 20a8b5713a2049acaf0c900d721dc2b25d147a5a) --- .gitignore | 8 + Dockerfile | 2 +- Dockerfile-clustertests | 13 +- Dockerfile-clustertests-client | 10 +- Dockerfile-datagen | 47 + Dockerfile-dax | 32 + Dockerfile-dax-quick | 18 + Makefile | 61 +- api.go | 518 ++- api_directive.go | 983 +++++ api_directive_internal_test.go | 25 + api_directive_test.go | 97 + batch/Makefile | 18 +- batch/README.md | 2 +- batch/batch.go | 4 + batch/batch_test.go | 4 +- batch/docker-compose.yml | 5 +- bsi.go | 32 +- bsi_test.go | 8 +- catcher.go | 2 +- client/api.go | 13 +- client/client.go | 108 +- client/client_it_test.go | 17 +- client/importer.go | 18 +- client/orm.go | 68 +- client/orm_test.go | 37 +- client/types/time.go | 49 + cluster.go | 101 +- cmd/cli.go | 8 + cmd/dax.go | 32 + cmd/import_test.go | 2 +- cmd/root.go | 8 +- ctl/cli.go | 321 +- ctl/dax.go | 39 + ctl/server.go | 242 +- dax/Makefile | 89 + dax/README.md | 64 + dax/address.go | 138 + dax/address_test.go | 176 + dax/boltdb/boltdb.go | 151 + dax/boltdb/boltdb_test.go | 17 + dax/boltdb/directiveversion.go | 66 + dax/boltdb/node.go | 157 + dax/boltdb/node_test.go | 55 + dax/boltdb/versionstore.go | 768 ++++ dax/boltdb/versionstore_test.go | 388 ++ dax/computer/alpha/alpha.go | 24 + dax/computer/alpha/snapshot.go | 104 + dax/computer/alpha/writelog.go | 331 ++ dax/computer/api/openapi.yaml | 235 ++ dax/computer/computer.go | 5 + dax/computer/snapshot.go | 61 + dax/computer/writelog.go | 307 ++ dax/dax.go | 11 + dax/directive.go | 175 + dax/docker-compose.yml | 66 + dax/errors.go | 80 + dax/fieldversion.go | 32 + dax/http/handler.go | 219 ++ dax/inmem/inmem.go | 2 + dax/inmem/versionstore.go | 430 ++ dax/inmem/versionstore_test.go | 166 + dax/mds/api/openapi.yaml | 548 +++ dax/mds/client/client.go | 478 +++ dax/mds/controller/alpha/director.go | 125 + dax/mds/controller/balancer.go | 69 + dax/mds/controller/config.go | 29 + dax/mds/controller/controller.go | 1892 +++++++++ dax/mds/controller/controller_test.go | 1373 +++++++ dax/mds/controller/director.go | 40 + dax/mds/controller/errors.go | 101 + dax/mds/controller/http/director.go | 185 + dax/mds/controller/naive/balancer.go | 539 +++ dax/mds/controller/naive/balancer_test.go | 739 ++++ dax/mds/controller/naive/boltdb/balancer.go | 515 +++ .../controller/naive/boltdb/balancer_test.go | 709 ++++ dax/mds/controller/naive/types.go | 81 + dax/mds/controller/node_registerer.go | 60 + dax/mds/controller/partitioner/partitioner.go | 55 + .../partitioner/partitioner_test.go | 86 + dax/mds/controller/sets.go | 169 + dax/mds/controller/stringers.go | 104 + dax/mds/controller/types.go | 19 + dax/mds/http/addressmanager.go | 76 + dax/mds/http/handler.go | 664 ++++ dax/mds/mds.go | 463 +++ dax/mds/poller/config.go | 15 + dax/mds/poller/interfaces.go | 59 + dax/mds/poller/poller.go | 156 + dax/mds/poller/poller_test.go | 182 + dax/mds/schemar/boltdb/schemar.go | 386 ++ dax/mds/schemar/boltdb/schemar_test.go | 146 + dax/mds/schemar/errors.go | 44 + dax/mds/schemar/http/handler.go | 199 + dax/mds/schemar/schemar.go | 54 + dax/mds/schemar/schemar_test.go | 124 + dax/node.go | 45 + dax/partition.go | 65 + dax/queryer/alpha/router.go | 36 + dax/queryer/api/openapi.yaml | 99 + dax/queryer/compute_api.go | 329 ++ dax/queryer/compute_importer.go | 68 + dax/queryer/config.go | 21 + dax/queryer/featurebase_importer.go | 41 + dax/queryer/http/handler.go | 113 + dax/queryer/importer.go | 129 + dax/queryer/interfaces.go | 78 + dax/queryer/orchestrator.go | 3494 +++++++++++++++++ dax/queryer/queryer.go | 304 ++ dax/queryer/router.go | 23 + dax/queryer/schema_api.go | 462 +++ dax/queryer/translator.go | 325 ++ dax/role.go | 57 + dax/server/config.go | 303 ++ dax/server/dup.go | 15 + dax/server/dup_arm64.go | 15 + dax/server/server.go | 523 +++ dax/shard.go | 64 + dax/snapshot.go | 34 + dax/snapshotter/api/openapi.yaml | 101 + dax/snapshotter/client/client.go | 109 + dax/snapshotter/config.go | 8 + dax/snapshotter/http/handler.go | 120 + dax/snapshotter/snapshotter.go | 112 + dax/snapshotter/snapshotter_test.go | 70 + dax/table.go | 581 +++ dax/table_test.go | 325 ++ dax/test/boltdb/helpers.go | 50 + dax/test/datagen/container.go | 80 + dax/test/dax/dax_test.go | 1916 +++++++++ dax/test/dax/image.go | 24 + dax/test/dax/testdata/basic.yaml | 52 + dax/test/docker/client.go | 34 + dax/test/docker/docker.go | 521 +++ dax/test/featurebase/container.go | 75 + dax/test/inspector/inspector.go | 217 + dax/test/schemar.go | 29 + dax/test/table.go | 69 + dax/time.go | 62 + dax/versionstore.go | 110 + dax/workerjob.go | 169 + dax/workerjob_test.go | 81 + dax/writelogger/api/openapi.yaml | 113 + dax/writelogger/client/client.go | 145 + dax/writelogger/config.go | 8 + dax/writelogger/http/handler.go | 121 + dax/writelogger/writelogger.go | 127 + dax/writelogger/writelogger_test.go | 73 + errors/errors.go | 6 +- errors/errors_test.go | 50 +- executor.go | 128 +- executor_internal_test.go | 4 +- fbcloud/auth.go | 79 + fbcloud/client.go | 147 + fragment.go | 6 +- go.mod | 40 +- go.sum | 66 +- handler.go | 6 + holder.go | 126 +- http_handler.go | 127 +- idk/Makefile | 22 +- idk/README.md | 12 +- idk/api/schema.go | 9 +- idk/datagen/cmd.go | 70 +- idk/datagen/custom.go | 11 +- idk/datagen/testdata/basic.yaml | 52 + idk/datagen/testdata/keyedtable.yaml | 58 + idk/datagen/testdata/keys_ids.yaml | 25 + idk/datagen/testdata/unkeyedtable.yaml | 69 + idk/docker-compose.yml | 37 +- idk/idktest/idktest.go | 5 + idk/ingest.go | 162 +- idk/ingest_test.go | 203 + idk/interfaces.go | 48 + idk/mds/importer.go | 276 ++ idk/mds/mds.go | 21 + idk/mds/schemamanager.go | 114 + idk/testdata/basic.yaml | 52 + idk/testdata/keyedtable.yaml | 58 + idk/testdata/unkeyedtable.yaml | 69 + index.go | 51 +- interfaces.go | 33 + internal/clustertests/cluster_test.go | 86 - internal/clustertests/docker-compose.yml | 12 +- internal_client.go | 104 +- logger/logger.go | 3 + rbf/cfg/cfg.go | 24 +- row.go | 199 +- server.go | 73 +- server/config.go | 32 + server/server.go | 210 +- server/server_test.go | 1 + sql.go | 9 + sql3/planner/executionplanner.go | 13 - sql3/planner/executionplannersystemtables.go | 7 + sql3/test/defs/defs_keyed.go | 131 +- sql3/test/defs/defs_unkeyed.go | 13 + sql3/test/defs/types.go | 21 + test/pilosa.go | 30 + test/transaction.go | 2 +- transaction_test.go | 6 +- txfactory.go | 10 +- 202 files changed, 32259 insertions(+), 867 deletions(-) create mode 100644 Dockerfile-datagen create mode 100644 Dockerfile-dax create mode 100644 Dockerfile-dax-quick create mode 100644 api_directive.go create mode 100644 api_directive_internal_test.go create mode 100644 api_directive_test.go create mode 100644 client/types/time.go create mode 100644 cmd/dax.go create mode 100644 ctl/dax.go create mode 100644 dax/Makefile create mode 100644 dax/README.md create mode 100644 dax/address.go create mode 100644 dax/address_test.go create mode 100644 dax/boltdb/boltdb.go create mode 100644 dax/boltdb/boltdb_test.go create mode 100644 dax/boltdb/directiveversion.go create mode 100644 dax/boltdb/node.go create mode 100644 dax/boltdb/node_test.go create mode 100644 dax/boltdb/versionstore.go create mode 100644 dax/boltdb/versionstore_test.go create mode 100644 dax/computer/alpha/alpha.go create mode 100644 dax/computer/alpha/snapshot.go create mode 100644 dax/computer/alpha/writelog.go create mode 100644 dax/computer/api/openapi.yaml create mode 100644 dax/computer/computer.go create mode 100644 dax/computer/snapshot.go create mode 100644 dax/computer/writelog.go create mode 100644 dax/dax.go create mode 100644 dax/directive.go create mode 100644 dax/docker-compose.yml create mode 100644 dax/errors.go create mode 100644 dax/fieldversion.go create mode 100644 dax/http/handler.go create mode 100644 dax/inmem/inmem.go create mode 100644 dax/inmem/versionstore.go create mode 100644 dax/inmem/versionstore_test.go create mode 100644 dax/mds/api/openapi.yaml create mode 100644 dax/mds/client/client.go create mode 100644 dax/mds/controller/alpha/director.go create mode 100644 dax/mds/controller/balancer.go create mode 100644 dax/mds/controller/config.go create mode 100644 dax/mds/controller/controller.go create mode 100644 dax/mds/controller/controller_test.go create mode 100644 dax/mds/controller/director.go create mode 100644 dax/mds/controller/errors.go create mode 100644 dax/mds/controller/http/director.go create mode 100644 dax/mds/controller/naive/balancer.go create mode 100644 dax/mds/controller/naive/balancer_test.go create mode 100644 dax/mds/controller/naive/boltdb/balancer.go create mode 100644 dax/mds/controller/naive/boltdb/balancer_test.go create mode 100644 dax/mds/controller/naive/types.go create mode 100644 dax/mds/controller/node_registerer.go create mode 100644 dax/mds/controller/partitioner/partitioner.go create mode 100644 dax/mds/controller/partitioner/partitioner_test.go create mode 100644 dax/mds/controller/sets.go create mode 100644 dax/mds/controller/stringers.go create mode 100644 dax/mds/controller/types.go create mode 100644 dax/mds/http/addressmanager.go create mode 100644 dax/mds/http/handler.go create mode 100644 dax/mds/mds.go create mode 100644 dax/mds/poller/config.go create mode 100644 dax/mds/poller/interfaces.go create mode 100644 dax/mds/poller/poller.go create mode 100644 dax/mds/poller/poller_test.go create mode 100644 dax/mds/schemar/boltdb/schemar.go create mode 100644 dax/mds/schemar/boltdb/schemar_test.go create mode 100644 dax/mds/schemar/errors.go create mode 100644 dax/mds/schemar/http/handler.go create mode 100644 dax/mds/schemar/schemar.go create mode 100644 dax/mds/schemar/schemar_test.go create mode 100644 dax/node.go create mode 100644 dax/partition.go create mode 100644 dax/queryer/alpha/router.go create mode 100644 dax/queryer/api/openapi.yaml create mode 100644 dax/queryer/compute_api.go create mode 100644 dax/queryer/compute_importer.go create mode 100644 dax/queryer/config.go create mode 100644 dax/queryer/featurebase_importer.go create mode 100644 dax/queryer/http/handler.go create mode 100644 dax/queryer/importer.go create mode 100644 dax/queryer/interfaces.go create mode 100644 dax/queryer/orchestrator.go create mode 100644 dax/queryer/queryer.go create mode 100644 dax/queryer/router.go create mode 100644 dax/queryer/schema_api.go create mode 100644 dax/queryer/translator.go create mode 100644 dax/role.go create mode 100644 dax/server/config.go create mode 100644 dax/server/dup.go create mode 100644 dax/server/dup_arm64.go create mode 100644 dax/server/server.go create mode 100644 dax/shard.go create mode 100644 dax/snapshot.go create mode 100644 dax/snapshotter/api/openapi.yaml create mode 100644 dax/snapshotter/client/client.go create mode 100644 dax/snapshotter/config.go create mode 100644 dax/snapshotter/http/handler.go create mode 100644 dax/snapshotter/snapshotter.go create mode 100644 dax/snapshotter/snapshotter_test.go create mode 100644 dax/table.go create mode 100644 dax/table_test.go create mode 100644 dax/test/boltdb/helpers.go create mode 100644 dax/test/datagen/container.go create mode 100644 dax/test/dax/dax_test.go create mode 100644 dax/test/dax/image.go create mode 100644 dax/test/dax/testdata/basic.yaml create mode 100644 dax/test/docker/client.go create mode 100644 dax/test/docker/docker.go create mode 100644 dax/test/featurebase/container.go create mode 100644 dax/test/inspector/inspector.go create mode 100644 dax/test/schemar.go create mode 100644 dax/test/table.go create mode 100644 dax/time.go create mode 100644 dax/versionstore.go create mode 100644 dax/workerjob.go create mode 100644 dax/workerjob_test.go create mode 100644 dax/writelogger/api/openapi.yaml create mode 100644 dax/writelogger/client/client.go create mode 100644 dax/writelogger/config.go create mode 100644 dax/writelogger/http/handler.go create mode 100644 dax/writelogger/writelogger.go create mode 100644 dax/writelogger/writelogger_test.go create mode 100644 fbcloud/auth.go create mode 100644 fbcloud/client.go create mode 100644 idk/datagen/testdata/basic.yaml create mode 100644 idk/datagen/testdata/keyedtable.yaml create mode 100644 idk/datagen/testdata/keys_ids.yaml create mode 100644 idk/datagen/testdata/unkeyedtable.yaml create mode 100644 idk/mds/importer.go create mode 100644 idk/mds/mds.go create mode 100644 idk/mds/schemamanager.go create mode 100644 idk/testdata/basic.yaml create mode 100644 idk/testdata/keyedtable.yaml create mode 100644 idk/testdata/unkeyedtable.yaml create mode 100644 interfaces.go create mode 100644 sql.go 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 }