Repository: coinbase/rosetta-cli
Pull request #118 (2020-08-23 21:06:59)

ZSTD Train Command
Blocking PR: https://github.com/coinbase/rosetta-sdk-go/pull/115

This PR adds a `utils:train-zstd` command that allows for training Zstandard dictionaries for small objects. There is more detail on this process in the linked PR.

### Changes
- [x] Allow training from an existing dictionary (i.e. decode from one dictionary encoding and convert to another)


Patch:
diff --git a/README.md b/README.md
index 38c9903..b3b7b99 100644
--- a/README.md
+++ b/README.md
@@ -47,6 +47,7 @@ Available Commands:
   configuration:validate       Ensure a configuration file at the provided path is formatted correctly
   help                         Help about any command
   utils:asserter-configuration Generate a static configuration file for the Asserter
+  utils:train-zstd             Generate a zstd dictionary for enhanced compression performance
   version                      Print rosetta-cli version
   view:account                 View an account balance
   view:block                   View a block
@@ -344,6 +345,35 @@ Global Flags:
                                     default values.
 ```
 
+#### utils:train-zstd
+```
+Zstandard (https://github.com/facebook/zstd) is used by
+rosetta-sdk-go/storage to compress data stored to disk. It is possible
+to improve compression performance by training a dictionary on a particular
+storage namespace. This command runs this training and outputs a dictionary
+that can be used with rosetta-sdk-go/storage.
+
+The arguments for this command are:
+<namespace> <database path> <dictionary path> <max items> (<existing dictionary path>)
+
+You can learn more about dictionary compression on the Zstandard
+website: https://github.com/facebook/zstd#the-case-for-small-data-compression
+
+Usage:
+  rosetta-cli utils:train-zstd [flags]
+
+Flags:
+  -h, --help   help for utils:train-zstd
+
+Global Flags:
+      --configuration-file string   Configuration file that provides connection and test settings.
+                                    If you would like to generate a starter configuration file (populated
+                                    with the defaults), run rosetta-cli configuration:create.
+
+                                    Any fields not populated in the configuration file will be populated with
+                                    default values.
+```
+
 ## Development
 * `make deps` to install dependencies
 * `make test` to run tests
diff --git a/cmd/root.go b/cmd/root.go
index 5a3cd06..bf8d0a2 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -84,6 +84,7 @@ default values.`,
 
 	// Utils
 	rootCmd.AddCommand(utilsAsserterConfigurationCmd)
+	rootCmd.AddCommand(utilsTrainZstdCmd)
 }
 
 func initConfig() {
diff --git a/cmd/utils_train_zstd.go b/cmd/utils_train_zstd.go
new file mode 100644
index 0000000..7ebce92
--- /dev/null
+++ b/cmd/utils_train_zstd.go
@@ -0,0 +1,78 @@
+// Copyright 2020 Coinbase, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package cmd
+
+import (
+	"context"
+	"log"
+	"path"
+	"strconv"
+
+	"github.com/coinbase/rosetta-sdk-go/storage"
+	"github.com/spf13/cobra"
+)
+
+const (
+	trainArgs = 4
+)
+
+var (
+	utilsTrainZstdCmd = &cobra.Command{
+		Use:   "utils:train-zstd",
+		Short: "Generate a zstd dictionary for enhanced compression performance",
+		Long: `Zstandard (https://github.com/facebook/zstd) is used by
+rosetta-sdk-go/storage to compress data stored to disk. It is possible
+to improve compression performance by training a dictionary on a particular
+storage namespace. This command runs this training and outputs a dictionary
+that can be used with rosetta-sdk-go/storage.
+
+The arguments for this command are:
+<namespace> <database path> <dictionary path> <max items> (<existing dictionary path>)
+
+You can learn more about dictionary compression on the Zstandard
+website: https://github.com/facebook/zstd#the-case-for-small-data-compression`,
+		Run:  runTrainZstdCmd,
+		Args: cobra.MinimumNArgs(trainArgs),
+	}
+)
+
+func runTrainZstdCmd(cmd *cobra.Command, args []string) {
+	ctx := context.Background()
+
+	namespace := args[0]
+	databasePath := path.Clean(args[1])
+	dictionaryPath := path.Clean(args[2])
+	maxItems, err := strconv.Atoi(args[3])
+	if err != nil {
+		log.Fatalf("%s: unable to convert max items to integer", err.Error())
+	}
+
+	compressorEntries := []*storage.CompressorEntry{}
+	if len(args) > trainArgs {
+		compressorEntries = append(compressorEntries, &storage.CompressorEntry{
+			Namespace:      namespace,
+			DictionaryPath: args[4],
+		})
+
+		log.Printf("found dictionary path %s\n", args[4])
+	}
+
+	log.Printf("Running zstd training (this could take a while)...")
+
+	_, _, err = storage.BadgerTrain(ctx, namespace, databasePath, dictionaryPath, maxItems, compressorEntries)
+	if err != nil {
+		log.Fatal(err)
+	}
+}
diff --git a/go.mod b/go.mod
index b872584..b827b4a 100644
--- a/go.mod
+++ b/go.mod
@@ -3,7 +3,7 @@ module github.com/coinbase/rosetta-cli
 go 1.13
 
 require (
-	github.com/coinbase/rosetta-sdk-go v0.3.5-0.20200821143659-1916d6b4bde3
+	github.com/coinbase/rosetta-sdk-go v0.3.5-0.20200824221853-d7b1fe2f9239
 	github.com/fatih/color v1.9.0
 	github.com/jinzhu/copier v0.0.0-20190924061706-b57f9002281a
 	github.com/olekukonko/tablewriter v0.0.2-0.20190409134802-7e037d187b0c
diff --git a/go.sum b/go.sum
index 4cfe6d3..1b91140 100644
--- a/go.sum
+++ b/go.sum
@@ -71,6 +71,26 @@ github.com/coinbase/rosetta-sdk-go v0.3.5-0.20200820205631-4c7341da957f h1:jkgbB
 github.com/coinbase/rosetta-sdk-go v0.3.5-0.20200820205631-4c7341da957f/go.mod h1:CuKG9s+UkPYUKgm13Ya3yxn/7qKTob26BKlbwdBwKno=
 github.com/coinbase/rosetta-sdk-go v0.3.5-0.20200821143659-1916d6b4bde3 h1:/XBnrmEdVH8j7dsDIArND3leI3JWOq2M1UVTGtBtmyQ=
 github.com/coinbase/rosetta-sdk-go v0.3.5-0.20200821143659-1916d6b4bde3/go.mod h1:CuKG9s+UkPYUKgm13Ya3yxn/7qKTob26BKlbwdBwKno=
+github.com/coinbase/rosetta-sdk-go v0.3.5-0.20200823205952-7cecd3cb6e87 h1:Kzg6yIEvndcxCmj0W8pwrDXjDR8T8uab14qwzFO7M8k=
+github.com/coinbase/rosetta-sdk-go v0.3.5-0.20200823205952-7cecd3cb6e87/go.mod h1:/dKD5dZJ0bzDXrd/M3XMQmpXRk0bru3s9EMIeyBJfy4=
+github.com/coinbase/rosetta-sdk-go v0.3.5-0.20200823213836-352196d0d456 h1:XmkCep09YaBQYcZuVlctAfWuF/fHBKe456UMpq4mNiM=
+github.com/coinbase/rosetta-sdk-go v0.3.5-0.20200823213836-352196d0d456/go.mod h1:/dKD5dZJ0bzDXrd/M3XMQmpXRk0bru3s9EMIeyBJfy4=
+github.com/coinbase/rosetta-sdk-go v0.3.5-0.20200823214948-e4483cf68f05 h1:4oV9FwCHfbS/7DNHFl4Z1QzgV4TcEBvvLKkg5SpT4FM=
+github.com/coinbase/rosetta-sdk-go v0.3.5-0.20200823214948-e4483cf68f05/go.mod h1:/dKD5dZJ0bzDXrd/M3XMQmpXRk0bru3s9EMIeyBJfy4=
+github.com/coinbase/rosetta-sdk-go v0.3.5-0.20200823223149-c0198457c35b h1:SFzcSwlJq48TxtN26psU2GCt05+dPp0SRjfgkcu1jio=
+github.com/coinbase/rosetta-sdk-go v0.3.5-0.20200823223149-c0198457c35b/go.mod h1:/dKD5dZJ0bzDXrd/M3XMQmpXRk0bru3s9EMIeyBJfy4=
+github.com/coinbase/rosetta-sdk-go v0.3.5-0.20200823224719-50905784782b h1:sQ3DvCS3BXgo4jfdQLl1gaB4CA7RjpswCTc8MXz17KM=
+github.com/coinbase/rosetta-sdk-go v0.3.5-0.20200823224719-50905784782b/go.mod h1:/dKD5dZJ0bzDXrd/M3XMQmpXRk0bru3s9EMIeyBJfy4=
+github.com/coinbase/rosetta-sdk-go v0.3.5-0.20200823231142-f9460bf368a6 h1:8xfgGrBrC2vxkkFAXuv/zfpU2CZi8kJpMTOPthtdngc=
+github.com/coinbase/rosetta-sdk-go v0.3.5-0.20200823231142-f9460bf368a6/go.mod h1:/dKD5dZJ0bzDXrd/M3XMQmpXRk0bru3s9EMIeyBJfy4=
+github.com/coinbase/rosetta-sdk-go v0.3.5-0.20200823232820-91097515172b h1:fz7eiGsWwhSlIkgubh4/kACZlpsf74wYra5yXt6P/VY=
+github.com/coinbase/rosetta-sdk-go v0.3.5-0.20200823232820-91097515172b/go.mod h1:/dKD5dZJ0bzDXrd/M3XMQmpXRk0bru3s9EMIeyBJfy4=
+github.com/coinbase/rosetta-sdk-go v0.3.5-0.20200824171955-1e3958442bc9 h1:iMAeF253Zdz7LgD0WUAqBnutSMCOItcF+DeX7VKS4T4=
+github.com/coinbase/rosetta-sdk-go v0.3.5-0.20200824171955-1e3958442bc9/go.mod h1:/dKD5dZJ0bzDXrd/M3XMQmpXRk0bru3s9EMIeyBJfy4=
+github.com/coinbase/rosetta-sdk-go v0.3.5-0.20200824172949-61aa33d621c1 h1:F9f2IlUT0KTssiEjxwv/37SmYdiPSboeNyCSsNuVa90=
+github.com/coinbase/rosetta-sdk-go v0.3.5-0.20200824172949-61aa33d621c1/go.mod h1:/dKD5dZJ0bzDXrd/M3XMQmpXRk0bru3s9EMIeyBJfy4=
+github.com/coinbase/rosetta-sdk-go v0.3.5-0.20200824221853-d7b1fe2f9239 h1:eP88DoIkcbMG5zHja3lUh4F2c+9am1T0zk/AChy1Vpo=
+github.com/coinbase/rosetta-sdk-go v0.3.5-0.20200824221853-d7b1fe2f9239/go.mod h1:/dKD5dZJ0bzDXrd/M3XMQmpXRk0bru3s9EMIeyBJfy4=
 github.com/coreos/bbolt v1.3.2/go.mod h1:iRUV2dpdMOn7Bo10OQBFzIJO9kkE559Wcmn+qkEiiKk=
 github.com/coreos/etcd v3.3.10+incompatible/go.mod h1:uF7uidLiAD3TWHmW31ZFd/JWoc32PjwdhPthX9715RE=
 github.com/coreos/go-etcd v2.0.0+incompatible/go.mod h1:Jez6KQU2B/sWsbdaef3ED8NzMklzPG4d5KIOhIy30Tk=


Test patch:
diff --git a/pkg/tester/construction.go b/pkg/tester/construction.go
index b026e32..891d0bd 100644
--- a/pkg/tester/construction.go
+++ b/pkg/tester/construction.go
@@ -67,7 +67,11 @@ func InitializeConstruction(
 		log.Fatalf("%s: cannot create command path", err.Error())
 	}
 
-	localStore, err := storage.NewBadgerStorage(ctx, dataPath, config.DisableMemoryLimit)
+	opts := []storage.BadgerOption{}
+	if !config.DisableMemoryLimit {
+		opts = append(opts, storage.WithMemoryLimit())
+	}
+	localStore, err := storage.NewBadgerStorage(ctx, dataPath, opts...)
 	if err != nil {
 		log.Fatalf("%s: unable to initialize database", err.Error())
 	}
diff --git a/pkg/tester/data.go b/pkg/tester/data.go
index 1670d30..8b397f2 100644
--- a/pkg/tester/data.go
+++ b/pkg/tester/data.go
@@ -139,7 +139,11 @@ func InitializeData(
 		log.Fatalf("%s: cannot create command path", err.Error())
 	}
 
-	localStore, err := storage.NewBadgerStorage(ctx, dataPath, config.DisableMemoryLimit)
+	opts := []storage.BadgerOption{}
+	if !config.DisableMemoryLimit {
+		opts = append(opts, storage.WithMemoryLimit())
+	}
+	localStore, err := storage.NewBadgerStorage(ctx, dataPath, opts...)
 	if err != nil {
 		log.Fatalf("%s: unable to initialize database", err.Error())
 	}
@@ -558,7 +562,11 @@ func (t *DataTester) recursiveOpSearch(
 	}
 	defer utils.RemoveTempDir(tmpDir)
 
-	localStore, err := storage.NewBadgerStorage(ctx, tmpDir, t.config.DisableMemoryLimit)
+	opts := []storage.BadgerOption{}
+	if !t.config.DisableMemoryLimit {
+		opts = append(opts, storage.WithMemoryLimit())
+	}
+	localStore, err := storage.NewBadgerStorage(ctx, tmpDir, opts...)
 	if err != nil {
 		return nil, fmt.Errorf("%w: unable to initialize database", err)
 	}
diff --git a/pkg/tester/data_results_test.go b/pkg/tester/data_results_test.go
index 25f768c..25950f4 100644
--- a/pkg/tester/data_results_test.go
+++ b/pkg/tester/data_results_test.go
@@ -305,7 +305,7 @@ func TestComputeCheckDataResults(t *testing.T) {
 				assert.NoError(t, err)
 
 				ctx := context.Background()
-				localStore, err := storage.NewBadgerStorage(ctx, dir, false)
+				localStore, err := storage.NewBadgerStorage(ctx, dir)
 				assert.NoError(t, err)
 
 				logPath := path.Join(dir, "results.json")