diff --git a/api.go b/api.go index eeb81f0c7..c40cc1a93 100644 --- a/api.go +++ b/api.go @@ -280,6 +280,15 @@ func (api *API) DeleteIndex(ctx context.Context, indexName string) error { return nil } +func (api *API) WriteColumnAttrDataTo(ctx context.Context, w io.Writer, indexName string) error { + index := api.holder.Index(indexName) + if index == nil { + return newNotFoundError(ErrIndexNotFound, indexName) + } + _, err := index.ColumnAttrStore().WriteTo(w) + return err +} + // CreateField makes the named field in the named index with the given options. // This method currently only takes a single functional option, but that may be // changed in the future to support multiple options. @@ -337,6 +346,15 @@ func (api *API) Field(ctx context.Context, indexName, fieldName string) (*Field, return field, nil } +func (api *API) WriteRowAttrDataTo(ctx context.Context, w io.Writer, indexName, fieldName string) error { + field := api.holder.Field(indexName, fieldName) + if field == nil { + return newNotFoundError(ErrFieldNotFound, fieldName) + } + _, err := field.RowAttrStore().WriteTo(w) + return err +} + func setUpImportOptions(opts ...ImportOption) (*ImportOptions, error) { options := &ImportOptions{} for _, opt := range opts { @@ -824,6 +842,34 @@ func (api *API) TranslateData(ctx context.Context, indexName string, partition i return store, nil } +// FieldTranslateData returns all translation data in the specified field. +func (api *API) FieldTranslateData(ctx context.Context, indexName, fieldName string) (io.WriterTo, error) { + span, _ := tracing.StartSpanFromContext(ctx, "API.FieldTranslateData") + defer span.Finish() + if err := api.validate(apiFieldTranslateData); err != nil { + return nil, errors.Wrap(err, "validating api method") + } + + // Retrieve index from holder. + idx := api.holder.Index(indexName) + if idx == nil { + return nil, newNotFoundError(ErrIndexNotFound, indexName) + } + + // Retrieve field from index. + field := idx.Field(fieldName) + if field == nil { + return nil, newNotFoundError(ErrFieldNotFound, fieldName) + } + + // Retrieve translatestore from holder. + store := field.TranslateStore() + if store == nil { + return nil, ErrTranslateStoreNotFound + } + return store, nil +} + // Hosts returns a list of the hosts in the cluster including their ID, // URL, and which is the primary. func (api *API) Hosts(ctx context.Context) []*topology.Node { @@ -1217,6 +1263,48 @@ func (api *API) FieldAttrDiff(ctx context.Context, indexName string, fieldName s return attrs, nil } +// IndexShardSnapshot returns a reader that contains the contents of an RBF snapshot for an index/shard. +func (api *API) IndexShardSnapshot(ctx context.Context, indexName string, shard uint64) (io.ReadCloser, error) { + span, _ := tracing.StartSpanFromContext(ctx, "API.IndexShardSnapshot") + defer span.Finish() + + // Find index. + index := api.holder.Index(indexName) + if index == nil { + return nil, newNotFoundError(ErrIndexNotFound, indexName) + } + + // Start transaction. + tx := index.holder.txf.NewTx(Txo{Index: index, Shard: shard}) + + // Ensure transaction is an RBF transaction. + rtx, ok := tx.(*RBFTx) + if !ok { + tx.Rollback() + return nil, fmt.Errorf("snapshot not available for %q storage", tx.Type()) + } + + r, err := rtx.SnapshotReader() + if err != nil { + tx.Rollback() + return nil, err + } + return &txReadCloser{tx: tx, Reader: r}, nil +} + +var _ io.ReadCloser = (*txReadCloser)(nil) + +// txReadCloser wraps a reader to close a tx on close. +type txReadCloser struct { + io.Reader + tx Tx +} + +func (r *txReadCloser) Close() error { + r.tx.Rollback() + return nil +} + // ImportOptions holds the options for the API.Import // method. // @@ -1756,6 +1844,19 @@ func (api *API) AvailableShardsByIndex(ctx context.Context) map[string]*roaring. return api.holder.availableShardsByIndex() } +// AvailableShards returns bitmap of available shards for a single index. +func (api *API) AvailableShards(ctx context.Context, indexName string) (*roaring.Bitmap, error) { + span, _ := tracing.StartSpanFromContext(ctx, "API.AvailableShards") + defer span.Finish() + + // Find the index. + index := api.holder.Index(indexName) + if index == nil { + return nil, newNotFoundError(ErrIndexNotFound, indexName) + } + return index.AvailableShards(false), nil +} + // StatsWithTags returns an instance of whatever implementation of StatsClient // pilosa is using with the given tags. func (api *API) StatsWithTags(tags []string) stats.StatsClient { @@ -2214,6 +2315,11 @@ func (api *API) ResetIDAlloc(index string) error { return api.holder.ida.reset(index) } +func (api *API) WriteIDAllocDataTo(w io.Writer) error { + _, err := api.holder.ida.WriteTo(w) + return err +} + // TranslateIndexDB is an internal function to load the index keys database // rd is a boltdb file. func (api *API) TranslateIndexDB(ctx context.Context, indexName string, partitionID int, rd io.Reader) error { @@ -2261,6 +2367,7 @@ const ( apiFragmentBlocks apiFragmentData apiTranslateData + apiFieldTranslateData apiField apiFieldAttrDiff //apiHosts // not implemented @@ -2299,10 +2406,11 @@ var methodsCommon = map[apiMethod]struct{}{ } var methodsResizing = map[apiMethod]struct{}{ - apiFragmentData: {}, - apiTranslateData: {}, - apiResizeAbort: {}, - apiSchema: {}, + apiFragmentData: {}, + apiTranslateData: {}, + apiFieldTranslateData: {}, + apiResizeAbort: {}, + apiSchema: {}, } var methodsDegraded = map[apiMethod]struct{}{ @@ -2337,6 +2445,7 @@ var methodsNormal = map[apiMethod]struct{}{ apiFragmentBlockData: {}, apiFragmentBlocks: {}, apiField: {}, + apiFieldTranslateData: {}, apiFieldAttrDiff: {}, apiImport: {}, apiImportValue: {}, @@ -2352,6 +2461,7 @@ var methodsNormal = map[apiMethod]struct{}{ apiStartTransaction: {}, apiFinishTransaction: {}, apiTransactions: {}, + apiTranslateData: {}, apiGetTransaction: {}, apiActiveQueries: {}, apiPastQueries: {}, diff --git a/apimethod_string.go b/apimethod_string.go index 8851ec725..59549438d 100644 --- a/apimethod_string.go +++ b/apimethod_string.go @@ -20,35 +20,36 @@ func _() { _ = x[apiFragmentBlocks-9] _ = x[apiFragmentData-10] _ = x[apiTranslateData-11] - _ = x[apiField-12] - _ = x[apiFieldAttrDiff-13] - _ = x[apiImport-14] - _ = x[apiImportValue-15] - _ = x[apiIndex-16] - _ = x[apiIndexAttrDiff-17] - _ = x[apiQuery-18] - _ = x[apiRecalculateCaches-19] - _ = x[apiRemoveNode-20] - _ = x[apiResizeAbort-21] - _ = x[apiSchema-22] - _ = x[apiShardNodes-23] - _ = x[apiState-24] - _ = x[apiViews-25] - _ = x[apiApplySchema-26] - _ = x[apiStartTransaction-27] - _ = x[apiFinishTransaction-28] - _ = x[apiTransactions-29] - _ = x[apiGetTransaction-30] - _ = x[apiActiveQueries-31] - _ = x[apiPastQueries-32] - _ = x[apiIDReserve-33] - _ = x[apiIDCommit-34] - _ = x[apiIDReset-35] + _ = x[apiFieldTranslateData-12] + _ = x[apiField-13] + _ = x[apiFieldAttrDiff-14] + _ = x[apiImport-15] + _ = x[apiImportValue-16] + _ = x[apiIndex-17] + _ = x[apiIndexAttrDiff-18] + _ = x[apiQuery-19] + _ = x[apiRecalculateCaches-20] + _ = x[apiRemoveNode-21] + _ = x[apiResizeAbort-22] + _ = x[apiSchema-23] + _ = x[apiShardNodes-24] + _ = x[apiState-25] + _ = x[apiViews-26] + _ = x[apiApplySchema-27] + _ = x[apiStartTransaction-28] + _ = x[apiFinishTransaction-29] + _ = x[apiTransactions-30] + _ = x[apiGetTransaction-31] + _ = x[apiActiveQueries-32] + _ = x[apiPastQueries-33] + _ = x[apiIDReserve-34] + _ = x[apiIDCommit-35] + _ = x[apiIDReset-36] } -const _apiMethod_name = "apiClusterMessageapiCreateFieldapiCreateIndexapiDeleteFieldapiDeleteAvailableShardapiDeleteIndexapiDeleteViewapiExportCSVapiFragmentBlockDataapiFragmentBlocksapiFragmentDataapiTranslateDataapiFieldapiFieldAttrDiffapiImportapiImportValueapiIndexapiIndexAttrDiffapiQueryapiRecalculateCachesapiRemoveNodeapiResizeAbortapiSchemaapiShardNodesapiStateapiViewsapiApplySchemaapiStartTransactionapiFinishTransactionapiTransactionsapiGetTransactionapiActiveQueriesapiPastQueriesapiIDReserveapiIDCommitapiIDReset" +const _apiMethod_name = "apiClusterMessageapiCreateFieldapiCreateIndexapiDeleteFieldapiDeleteAvailableShardapiDeleteIndexapiDeleteViewapiExportCSVapiFragmentBlockDataapiFragmentBlocksapiFragmentDataapiTranslateDataapiFieldTranslateDataapiFieldapiFieldAttrDiffapiImportapiImportValueapiIndexapiIndexAttrDiffapiQueryapiRecalculateCachesapiRemoveNodeapiResizeAbortapiSchemaapiShardNodesapiStateapiViewsapiApplySchemaapiStartTransactionapiFinishTransactionapiTransactionsapiGetTransactionapiActiveQueriesapiPastQueriesapiIDReserveapiIDCommitapiIDReset" -var _apiMethod_index = [...]uint16{0, 17, 31, 45, 59, 82, 96, 109, 121, 141, 158, 173, 189, 197, 213, 222, 236, 244, 260, 268, 288, 301, 315, 324, 337, 345, 353, 367, 386, 406, 421, 438, 454, 468, 480, 491, 501} +var _apiMethod_index = [...]uint16{0, 17, 31, 45, 59, 82, 96, 109, 121, 141, 158, 173, 189, 210, 218, 234, 243, 257, 265, 281, 289, 309, 322, 336, 345, 358, 366, 374, 388, 407, 427, 442, 459, 475, 489, 501, 512, 522} func (i apiMethod) String() string { if i < 0 || i >= apiMethod(len(_apiMethod_index)-1) { diff --git a/attr.go b/attr.go index 66667f806..fa5d11eca 100644 --- a/attr.go +++ b/attr.go @@ -16,6 +16,7 @@ package pilosa import ( "bytes" + "io" "sort" "github.com/gogo/protobuf/proto" @@ -32,6 +33,8 @@ const ( // AttrStore represents an interface for handling row/column attributes. type AttrStore interface { + io.WriterTo + Path() string Open() error Close() error @@ -76,6 +79,9 @@ func (s nopAttrStore) Blocks() ([]AttrBlock, error) { return nil, nil } // BlockData is a no-op implementation of AttrStore BlockData method. func (s nopAttrStore) BlockData(i uint64) (map[uint64]map[string]interface{}, error) { return nil, nil } +// WriteTo is a no-op implementation of AttrStore WriteTo method. +func (s nopAttrStore) WriteTo(w io.Writer) (int64, error) { return 0, nil } + // AttrBlock represents a checksummed block of the attribute store. type AttrBlock struct { ID uint64 `json:"id"` diff --git a/boltdb/attrstore.go b/boltdb/attrstore.go index cf5c62473..7367e31f1 100644 --- a/boltdb/attrstore.go +++ b/boltdb/attrstore.go @@ -16,9 +16,9 @@ package boltdb import ( "bytes" - "encoding/binary" "fmt" + "io" "sort" "sync" "time" @@ -276,6 +276,16 @@ func (s *attrStore) BlockData(i uint64) (m map[uint64]map[string]interface{}, er return m, nil } +// WriteTo writes the underlying database to w. +func (s *attrStore) WriteTo(w io.Writer) (int64, error) { + tx, err := s.db.Begin(false) + if err != nil { + return 0, err + } + defer func() { _ = tx.Rollback() }() + return tx.WriteTo(w) +} + // txAttrs returns a map of attributes for an id. func txAttrs(tx *bolt.Tx, id uint64) (map[string]interface{}, error) { v := tx.Bucket([]byte("attrs")).Get(u64tob(id)) diff --git a/client.go b/client.go index 7da57a9f7..9985e1342 100644 --- a/client.go +++ b/client.go @@ -53,6 +53,7 @@ type FieldValue struct { type InternalClient interface { InternalQueryClient + AvailableShards(ctx context.Context, indexName string) ([]uint64, error) MaxShardByIndex(ctx context.Context) (map[string]uint64, error) Schema(ctx context.Context) ([]*IndexInfo, error) PostSchema(ctx context.Context, uri *pnet.URI, s *Schema, remote bool) error @@ -80,6 +81,13 @@ type InternalClient interface { RetrieveTranslatePartitionFromURI(ctx context.Context, index string, partition int, uri pnet.URI) (io.ReadCloser, error) ImportRoaring(ctx context.Context, uri *pnet.URI, index, field string, shard uint64, remote bool, req *ImportRoaringRequest) error ImportColumnAttrs(ctx context.Context, uri *pnet.URI, index string, req *ImportColumnAttrsRequest) error + ShardReader(ctx context.Context, index string, shard uint64) (io.ReadCloser, error) + + IDAllocDataReader(ctx context.Context) (io.ReadCloser, error) + IndexTranslateDataReader(ctx context.Context, index string, partitionID int) (io.ReadCloser, error) + IndexAttrDataReader(ctx context.Context, index string) (io.ReadCloser, error) + FieldTranslateDataReader(ctx context.Context, index, field string) (io.ReadCloser, error) + FieldAttrDataReader(ctx context.Context, index, field string) (io.ReadCloser, error) StartTransaction(ctx context.Context, id string, timeout time.Duration, exclusive bool) (*Transaction, error) FinishTransaction(ctx context.Context, id string) (*Transaction, error) @@ -159,6 +167,10 @@ func newNopInternalClient() nopInternalClient { var _ InternalClient = newNopInternalClient() +func (n nopInternalClient) AvailableShards(ctx context.Context, indexName string) ([]uint64, error) { + return nil, nil +} + func (n nopInternalClient) MaxShardByIndex(context.Context) (map[string]uint64, error) { return nil, nil } @@ -197,6 +209,30 @@ func (n nopInternalClient) ImportColumnAttrs(ctx context.Context, uri *pnet.URI, return nil } +func (n nopInternalClient) ShardReader(ctx context.Context, index string, shard uint64) (io.ReadCloser, error) { + return nil, nil +} + +func (n nopInternalClient) IDAllocDataReader(ctx context.Context) (io.ReadCloser, error) { + return nil, nil +} + +func (n nopInternalClient) IndexTranslateDataReader(ctx context.Context, index string, partitionID int) (io.ReadCloser, error) { + return nil, nil +} + +func (n nopInternalClient) IndexAttrDataReader(ctx context.Context, index string) (io.ReadCloser, error) { + return nil, nil +} + +func (n nopInternalClient) FieldTranslateDataReader(ctx context.Context, index, field string) (io.ReadCloser, error) { + return nil, nil +} + +func (n nopInternalClient) FieldAttrDataReader(ctx context.Context, index, field string) (io.ReadCloser, error) { + return nil, nil +} + func (n nopInternalClient) EnsureIndex(ctx context.Context, name string, options IndexOptions) error { return nil } diff --git a/cmd/backup.go b/cmd/backup.go new file mode 100644 index 000000000..8e272c526 --- /dev/null +++ b/cmd/backup.go @@ -0,0 +1,43 @@ +// Copyright 2017 Pilosa Corp. +// +// 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" + "io" + + "github.com/pilosa/pilosa/v2/ctl" + "github.com/spf13/cobra" +) + +func newBackupCommand(stdin io.Reader, stdout io.Writer, stderr io.Writer) *cobra.Command { + cmd := ctl.NewBackupCommand(stdin, stdout, stderr) + ccmd := &cobra.Command{ + Use: "backup", + Short: "Back up pilosa server", + Long: ` +Backs up a pilosa server to a local snapshot file. +`, + RunE: func(c *cobra.Command, args []string) error { + return cmd.Run(context.Background()) + }, + } + + flags := ccmd.Flags() + flags.StringVarP(&cmd.OutputPath, "output", "o", "", "output path to write to") + flags.StringVar(&cmd.Host, "host", "localhost:10101", "host:port of Pilosa.") + ctl.SetTLSConfig(flags, "", &cmd.TLS.CertificatePath, &cmd.TLS.CertificateKeyPath, &cmd.TLS.CACertPath, &cmd.TLS.SkipVerify, &cmd.TLS.EnableClientVerification) + return ccmd +} diff --git a/cmd/root.go b/cmd/root.go index 89f59bcfb..a6de7527e 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -63,6 +63,7 @@ at https://www.pilosa.com/docs/. _ = rc.PersistentFlags().MarkHidden("dry-run") rc.PersistentFlags().StringP("config", "c", "", "Configuration file to read from.") + rc.AddCommand(newBackupCommand(stdin, stdout, stderr)) rc.AddCommand(newCheckCommand(stdin, stdout, stderr)) rc.AddCommand(newConfigCommand(stdin, stdout, stderr)) rc.AddCommand(newExportCommand(stdin, stdout, stderr)) diff --git a/ctl/backup.go b/ctl/backup.go new file mode 100644 index 000000000..d92316b2f --- /dev/null +++ b/ctl/backup.go @@ -0,0 +1,391 @@ +// Copyright 2017 Pilosa Corp. +// +// 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 ctl + +import ( + "archive/tar" + "bytes" + "compress/gzip" + "context" + "encoding/json" + "fmt" + "io" + "os" + "path" + "path/filepath" + "time" + + "github.com/pilosa/pilosa/v2" + "github.com/pilosa/pilosa/v2/server" + "github.com/pilosa/pilosa/v2/topology" +) + +// BackupCommand represents a command for backing up a Pilosa node. +type BackupCommand struct { // nolint: maligned + // Destination host and port. + Host string `json:"host"` + + // Path to write the backup to. + OutputPath string + + // Reusable client. + client pilosa.InternalClient + + // Standard input/output + *pilosa.CmdIO + + TLS server.TLSConfig +} + +// NewBackupCommand returns a new instance of BackupCommand. +func NewBackupCommand(stdin io.Reader, stdout, stderr io.Writer) *BackupCommand { + return &BackupCommand{ + CmdIO: pilosa.NewCmdIO(stdin, stdout, stderr), + } +} + +// TempPath returns the path to the temporary file to write the archive to. +func (cmd *BackupCommand) TempPath() string { + dir, base := filepath.Split(cmd.OutputPath) + return filepath.Join(dir, "."+base) +} + +// Run executes the main program execution. +func (cmd *BackupCommand) Run(ctx context.Context) error { + logger := cmd.Logger() + + // Validate arguments. + if cmd.OutputPath == "" { + return fmt.Errorf("-o flag required") + } + + // Create a client to the server. + client, err := commandClient(cmd) + if err != nil { + return fmt.Errorf("creating client: %w", err) + } + cmd.client = client + + // Determine the field type in order to correctly handle the input data. + indexes, err := cmd.client.Schema(ctx) + if err != nil { + return fmt.Errorf("getting schema: %w", err) + } + schema := &pilosa.Schema{Indexes: indexes} + + // Create output file in temporary location. + w, err := os.Create(cmd.OutputPath + ".tmp") + if err != nil { + return err + } + defer w.Close() + + // Open a tar/gzip writer to the temporary file. + gw := gzip.NewWriter(w) + defer gw.Close() + tw := tar.NewWriter(gw) + defer tw.Close() + + // Backup schema. + if err := cmd.backupSchema(ctx, tw, schema); err != nil { + return fmt.Errorf("cannot back up schema: %w", err) + } else if err := cmd.backupIDAllocData(ctx, tw); err != nil { + return fmt.Errorf("cannot back up id alloc data: %w", err) + } + + // Backup data for each index. + for _, ii := range schema.Indexes { + if err := cmd.backupIndex(ctx, tw, ii); err != nil { + return err + } + } + + // Move data file to final location. + logger.Printf("writing backup: %s", cmd.OutputPath) + if err := os.Rename(cmd.OutputPath+".tmp", cmd.OutputPath); err != nil { + return err + } + + return nil +} + +// backupSchema writes the schema to the archive. +func (cmd *BackupCommand) backupSchema(ctx context.Context, tw *tar.Writer, schema *pilosa.Schema) error { + logger := cmd.Logger() + logger.Printf("backing up schema") + + buf, err := json.MarshalIndent(schema, "", "\t") + if err != nil { + return fmt.Errorf("marshaling schema: %w", err) + } + + // Build header & copy data to archive. + if err = tw.WriteHeader(&tar.Header{ + Name: "schema", + Mode: 0666, + Size: int64(len(buf)), + ModTime: time.Now(), + }); err != nil { + return err + } else if _, err := tw.Write(buf); err != nil { + return fmt.Errorf("copying schema to archive: %w", err) + } + + return nil +} + +func (cmd *BackupCommand) backupIDAllocData(ctx context.Context, tw *tar.Writer) error { + logger := cmd.Logger() + logger.Printf("backing up id alloc data") + + rc, err := cmd.client.IDAllocDataReader(ctx) + if err != nil { + return fmt.Errorf("fetching id alloc data reader: %w", err) + } + defer rc.Close() + + // Read to buffer to determine size. + var buf bytes.Buffer + if _, err := buf.ReadFrom(rc); err != nil { + return fmt.Errorf("copying id alloc data to memory: %w", err) + } + + // Build header & copy data to archive. + if err = tw.WriteHeader(&tar.Header{ + Name: "idalloc", + Mode: 0666, + Size: int64(buf.Len()), + ModTime: time.Now(), + }); err != nil { + return err + } else if _, err := io.Copy(tw, &buf); err != nil { + return fmt.Errorf("copying id alloc data to archive: %w", err) + } + + return nil +} + +// backupIndex backs up all shards for a given index. +func (cmd *BackupCommand) backupIndex(ctx context.Context, tw *tar.Writer, ii *pilosa.IndexInfo) error { + logger := cmd.Logger() + logger.Printf("backing up index: %q", ii.Name) + + shards, err := cmd.client.AvailableShards(ctx, ii.Name) + if err != nil { + return fmt.Errorf("cannot find available shards for index %q: %w", ii.Name, err) + } + + // Back up all bitmap data for the index. + for _, shard := range shards { + if err := cmd.backupShard(ctx, tw, ii.Name, shard); err != nil { + return fmt.Errorf("cannot backup shard %d on index %q: %w", shard, ii.Name, err) + } + } + + // Back up translation data after bitmap data so we ensure we can translate all data. + if err := cmd.backupIndexTranslateData(ctx, tw, ii.Name); err != nil { + return err + } + if err := cmd.backupIndexAttrData(ctx, tw, ii.Name); err != nil { + return err + } + + // Back up field translation & attribute data. + for _, fi := range ii.Fields { + if err := cmd.backupFieldTranslateData(ctx, tw, ii.Name, fi.Name); err != nil { + return fmt.Errorf("cannot backup field translation data for field %q on index %q: %w", fi.Name, ii.Name, err) + } + if err := cmd.backupFieldAttrData(ctx, tw, ii.Name, fi.Name); err != nil { + return fmt.Errorf("cannot backup field attr data for field %q on index %q: %w", fi.Name, ii.Name, err) + } + } + + return nil +} + +// backupShard backs up a single shard from a single index. +func (cmd *BackupCommand) backupShard(ctx context.Context, tw *tar.Writer, indexName string, shard uint64) error { + logger := cmd.Logger() + logger.Printf("backing up shard: index=%q id=%d", indexName, shard) + + filename := path.Join("indexes", indexName, "shards", fmt.Sprintf("%04d", shard)) + + rc, err := cmd.client.ShardReader(ctx, indexName, shard) + if err != nil { + return fmt.Errorf("fetching shard reader: %w", err) + } + defer rc.Close() + + // Read to buffer to determine size. + // TODO: Provide size via the reader itself. + var buf bytes.Buffer + if _, err := buf.ReadFrom(rc); err != nil { + return fmt.Errorf("copying shard data to memory: %w", err) + } + + // Build header & copy data to archive. + if err = tw.WriteHeader(&tar.Header{ + Name: filename, + Mode: 0666, + Size: int64(buf.Len()), + ModTime: time.Now(), + }); err != nil { + return err + } else if _, err := io.Copy(tw, &buf); err != nil { + return fmt.Errorf("copying shard data to archive: %w", err) + } + + return nil +} + +func (cmd *BackupCommand) backupIndexTranslateData(ctx context.Context, tw *tar.Writer, name string) error { + // TODO: Fetch holder partition count. + partitionN := topology.DefaultPartitionN + for partitionID := 0; partitionID < partitionN; partitionID++ { + if err := cmd.backupIndexPartitionTranslateData(ctx, tw, name, partitionID); err != nil { + return fmt.Errorf("cannot backup index translation data for partition %d on %q: %w", partitionID, name, err) + } + } + return nil +} + +func (cmd *BackupCommand) backupIndexPartitionTranslateData(ctx context.Context, tw *tar.Writer, name string, partitionID int) error { + logger := cmd.Logger() + logger.Printf("backing up index translation data: %s/%d", name, partitionID) + + rc, err := cmd.client.IndexTranslateDataReader(ctx, name, partitionID) + if err == pilosa.ErrTranslateStoreNotFound { + return nil + } else if err != nil { + return fmt.Errorf("fetching translate data reader: %w", err) + } + defer rc.Close() + + // Read to buffer to determine size. + var buf bytes.Buffer + if _, err := buf.ReadFrom(rc); err != nil { + return fmt.Errorf("copying translate data to memory: %w", err) + } + + // Build header & copy data to archive. + if err = tw.WriteHeader(&tar.Header{ + Name: path.Join("indexes", name, "translate", fmt.Sprintf("%04d", partitionID)), + Mode: 0666, + Size: int64(buf.Len()), + ModTime: time.Now(), + }); err != nil { + return err + } else if _, err := io.Copy(tw, &buf); err != nil { + return fmt.Errorf("copying translate data to archive: %w", err) + } + + return nil +} + +func (cmd *BackupCommand) backupIndexAttrData(ctx context.Context, tw *tar.Writer, name string) error { + logger := cmd.Logger() + logger.Printf("backing up index attr data: %s", name) + + rc, err := cmd.client.IndexAttrDataReader(ctx, name) + if err != nil { + return fmt.Errorf("fetching index attr data reader: %w", err) + } + defer rc.Close() + + // Read to buffer to determine size. + var buf bytes.Buffer + if _, err := buf.ReadFrom(rc); err != nil { + return fmt.Errorf("copying index attr data to memory: %w", err) + } + + // Build header & copy data to archive. + if err = tw.WriteHeader(&tar.Header{ + Name: path.Join("indexes", name, "attributes"), + Mode: 0666, + Size: int64(buf.Len()), + ModTime: time.Now(), + }); err != nil { + return err + } else if _, err := io.Copy(tw, &buf); err != nil { + return fmt.Errorf("copying index attr data to archive: %w", err) + } + return nil +} + +func (cmd *BackupCommand) backupFieldTranslateData(ctx context.Context, tw *tar.Writer, indexName, fieldName string) error { + logger := cmd.Logger() + logger.Printf("backing up field translation data: %s/%s", indexName, fieldName) + + rc, err := cmd.client.FieldTranslateDataReader(ctx, indexName, fieldName) + if err == pilosa.ErrTranslateStoreNotFound { + return nil + } else if err != nil { + return fmt.Errorf("fetching translate data reader: %w", err) + } + defer rc.Close() + + // Read to buffer to determine size. + var buf bytes.Buffer + if _, err := buf.ReadFrom(rc); err != nil { + return fmt.Errorf("copying translate data to memory: %w", err) + } + + // Build header & copy data to archive. + if err = tw.WriteHeader(&tar.Header{ + Name: path.Join("indexes", indexName, "fields", fieldName, "translate"), + Mode: 0666, + Size: int64(buf.Len()), + ModTime: time.Now(), + }); err != nil { + return err + } else if _, err := io.Copy(tw, &buf); err != nil { + return fmt.Errorf("copying translate data to archive: %w", err) + } + return nil +} + +func (cmd *BackupCommand) backupFieldAttrData(ctx context.Context, tw *tar.Writer, indexName, fieldName string) error { + logger := cmd.Logger() + logger.Printf("backing up field attr data: %s/%s", indexName, fieldName) + + rc, err := cmd.client.FieldAttrDataReader(ctx, indexName, fieldName) + if err != nil { + return fmt.Errorf("fetching field attr data reader: %w", err) + } + defer rc.Close() + + // Read to buffer to determine size. + var buf bytes.Buffer + if _, err := buf.ReadFrom(rc); err != nil { + return fmt.Errorf("copying field attr data to memory: %w", err) + } + + // Build header & copy data to archive. + if err = tw.WriteHeader(&tar.Header{ + Name: path.Join("indexes", indexName, "fields", fieldName, "attributes"), + Mode: 0666, + Size: int64(buf.Len()), + ModTime: time.Now(), + }); err != nil { + return err + } else if _, err := io.Copy(tw, &buf); err != nil { + return fmt.Errorf("copying field attr data to archive: %w", err) + } + return nil +} + +func (cmd *BackupCommand) TLSHost() string { return cmd.Host } + +func (cmd *BackupCommand) TLSConfiguration() server.TLSConfig { return cmd.TLS } diff --git a/http/client.go b/http/client.go index 3adce376a..fdf2cf7d7 100644 --- a/http/client.go +++ b/http/client.go @@ -24,6 +24,7 @@ import ( "math/rand" "net/http" "net/url" + "path" "sort" "strconv" "time" @@ -104,6 +105,37 @@ func (c *InternalClient) maxShardByIndex(ctx context.Context) (map[string]uint64 return rsp.Standard, nil } +// AvailableShards returns a list of shards for an index. +func (c *InternalClient) AvailableShards(ctx context.Context, indexName string) ([]uint64, error) { + span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.AvailableShards") + defer span.Finish() + + // Execute request against the host. + u := uriPathToURL(c.defaultURI, path.Join("/internal/index", indexName, "/shards")) + + // Build request. + req, err := http.NewRequest("GET", u.String(), nil) + if err != nil { + return nil, errors.Wrap(err, "creating request") + } + + req.Header.Set("User-Agent", "pilosa/"+pilosa.Version) + req.Header.Set("Accept", "application/json") + + // Execute request. + resp, err := c.executeRequest(req.WithContext(ctx)) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + var rsp getIndexAvailableShardsResponse + if err := json.NewDecoder(resp.Body).Decode(&rsp); err != nil { + return nil, fmt.Errorf("json decode: %s", err) + } + return rsp.Shards, nil +} + // SchemaNode returns all index and field schema information from the specified // node. func (c *InternalClient) SchemaNode(ctx context.Context, uri *pnet.URI, views bool) ([]*pilosa.IndexInfo, error) { @@ -2079,6 +2111,157 @@ func (c *InternalClient) ImportFieldKeys(ctx context.Context, uri *pnet.URI, ind return nil } +// ShardReader returns a reader that provides a snapshot of the current shard RBF data. +func (c *InternalClient) ShardReader(ctx context.Context, index string, shard uint64) (io.ReadCloser, error) { + span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.ShardReader") + defer span.Finish() + + // Execute request against the host. + u := fmt.Sprintf("%s/internal/index/%s/shard/%d/snapshot", c.defaultURI, index, shard) + + // Build request. + req, err := http.NewRequest("GET", u, nil) + if err != nil { + return nil, errors.Wrap(err, "creating request") + } + + req.Header.Set("User-Agent", "pilosa/"+pilosa.Version) + req.Header.Set("Accept", "application/octet-stream") + + // Execute request. + resp, err := c.executeRequest(req.WithContext(ctx)) + if err != nil { + return nil, err + } + return resp.Body, nil +} + +// IDAllocDataReader returns a reader that provides a snapshot of ID allocation data. +func (c *InternalClient) IDAllocDataReader(ctx context.Context) (io.ReadCloser, error) { + span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.IDAllocDataReader") + defer span.Finish() + + // Build request. + req, err := http.NewRequest("GET", c.defaultURI.String()+"/internal/idalloc/data", nil) + if err != nil { + return nil, errors.Wrap(err, "creating request") + } + + req.Header.Set("User-Agent", "pilosa/"+pilosa.Version) + req.Header.Set("Accept", "application/octet-stream") + + // Execute request. + resp, err := c.executeRequest(req.WithContext(ctx)) + if err != nil { + return nil, err + } + return resp.Body, nil +} + +// IndexTranslateDataReader returns a reader that provides a snapshot of +// translation data for a partition in an index. +func (c *InternalClient) IndexTranslateDataReader(ctx context.Context, index string, partitionID int) (io.ReadCloser, error) { + span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.IndexTranslateDataReader") + 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) + + // Build request. + req, err := http.NewRequest("GET", u, nil) + if err != nil { + return nil, errors.Wrap(err, "creating request") + } + + req.Header.Set("User-Agent", "pilosa/"+pilosa.Version) + req.Header.Set("Accept", "application/octet-stream") + + // Execute request. + resp, err := c.executeRequest(req.WithContext(ctx)) + if resp != nil && resp.StatusCode == http.StatusNotFound { + resp.Body.Close() + return nil, pilosa.ErrTranslateStoreNotFound + } else if err != nil { + return nil, err + } + return resp.Body, nil +} + +// IndexAttrDataReader returns a reader that provides a snapshot of column attributes data. +func (c *InternalClient) IndexAttrDataReader(ctx context.Context, index string) (io.ReadCloser, error) { + span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.IndexAttrDataReader") + defer span.Finish() + + // Build request. + u := fmt.Sprintf("%s/internal/index/%s/attr/data", c.defaultURI.String(), url.QueryEscape(index)) + req, err := http.NewRequest("GET", u, nil) + if err != nil { + return nil, errors.Wrap(err, "creating request") + } + + req.Header.Set("User-Agent", "pilosa/"+pilosa.Version) + req.Header.Set("Accept", "application/octet-stream") + + // Execute request. + resp, err := c.executeRequest(req.WithContext(ctx)) + if err != nil { + return nil, err + } + return resp.Body, nil +} + +// FieldTranslateDataReader returns a reader that provides a snapshot of +// translation data for a field. +func (c *InternalClient) FieldTranslateDataReader(ctx context.Context, index, field string) (io.ReadCloser, error) { + span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.FieldTranslateDataReader") + 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)) + + // Build request. + req, err := http.NewRequest("GET", u, nil) + if err != nil { + return nil, errors.Wrap(err, "creating request") + } + + req.Header.Set("User-Agent", "pilosa/"+pilosa.Version) + req.Header.Set("Accept", "application/octet-stream") + + // Execute request. + resp, err := c.executeRequest(req.WithContext(ctx)) + if resp != nil && resp.StatusCode == http.StatusNotFound { + resp.Body.Close() + return nil, pilosa.ErrTranslateStoreNotFound + } else if err != nil { + return nil, err + } + return resp.Body, nil +} + +// FieldAttrDataReader returns a reader that provides a snapshot of row attributes data. +func (c *InternalClient) FieldAttrDataReader(ctx context.Context, index, field string) (io.ReadCloser, error) { + span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.FieldAttrDataReader") + defer span.Finish() + + // Build request. + u := fmt.Sprintf("%s/internal/index/%s/field/%s/attr/data", c.defaultURI.String(), url.QueryEscape(index), url.QueryEscape(field)) + req, err := http.NewRequest("GET", u, nil) + if err != nil { + return nil, errors.Wrap(err, "creating request") + } + + req.Header.Set("User-Agent", "pilosa/"+pilosa.Version) + req.Header.Set("Accept", "application/octet-stream") + + // Execute request. + resp, err := c.executeRequest(req.WithContext(ctx)) + if err != nil { + return nil, err + } + return resp.Body, nil +} + // Status function is just a public function for this particular implementation of InternalClient. // It's not require by pilosa.InternalClient interface. // The function returns pilosa cluster state as a string ("NORMAL", "DEGRADED", "DOWN", "RESIZING", ...) diff --git a/http/handler.go b/http/handler.go index ec976420d..94f14a011 100644 --- a/http/handler.go +++ b/http/handler.go @@ -44,6 +44,7 @@ import ( "github.com/pilosa/pilosa/v2/encoding/proto" "github.com/pilosa/pilosa/v2/logger" "github.com/pilosa/pilosa/v2/pql" + "github.com/pilosa/pilosa/v2/rbf" "github.com/pilosa/pilosa/v2/topology" "github.com/pilosa/pilosa/v2/tracing" "github.com/pkg/errors" @@ -229,7 +230,7 @@ func (h *Handler) populateValidators() { h.validators["GetIndex"] = queryValidationSpecRequired() h.validators["PostIndex"] = queryValidationSpecRequired() h.validators["DeleteIndex"] = queryValidationSpecRequired() - h.validators["GetTranslateData"] = queryValidationSpecRequired("index", "partition") + h.validators["GetTranslateData"] = queryValidationSpecRequired("index").Optional("partition", "field") h.validators["PostTranslateKeys"] = queryValidationSpecRequired() h.validators["PostField"] = queryValidationSpecRequired() h.validators["DeleteField"] = queryValidationSpecRequired() @@ -249,7 +250,9 @@ func (h *Handler) populateValidators() { h.validators["GetFragmentData"] = queryValidationSpecRequired("index", "field", "view", "shard") h.validators["GetFragmentNodes"] = queryValidationSpecRequired("shard", "index") h.validators["PostIndexAttrDiff"] = queryValidationSpecRequired() + h.validators["GetIndexAttrData"] = queryValidationSpecRequired() h.validators["PostFieldAttrDiff"] = queryValidationSpecRequired() + h.validators["GetFieldAttrData"] = queryValidationSpecRequired() h.validators["GetNodes"] = queryValidationSpecRequired() h.validators["GetShardMax"] = queryValidationSpecRequired() h.validators["GetTransactionList"] = queryValidationSpecRequired() @@ -424,12 +427,16 @@ func newRouter(handler *Handler) http.Handler { router.HandleFunc("/internal/fragment/data", handler.handleGetFragmentData).Methods("GET").Name("GetFragmentData") router.HandleFunc("/internal/fragment/nodes", handler.handleGetFragmentNodes).Methods("GET").Name("GetFragmentNodes") router.HandleFunc("/internal/index/{index}/attr/diff", handler.handlePostIndexAttrDiff).Methods("POST").Name("PostIndexAttrDiff") + router.HandleFunc("/internal/index/{index}/attr/data", handler.handleGetIndexAttrData).Methods("GET").Name("GetIndexAttrData") router.HandleFunc("/internal/translate/data", handler.handleGetTranslateData).Methods("GET").Name("GetTranslateData") router.HandleFunc("/internal/translate/data", handler.handlePostTranslateData).Methods("POST").Name("PostTranslateData") router.HandleFunc("/internal/translate/keys", handler.handlePostTranslateKeys).Methods("POST").Name("PostTranslateKeys") router.HandleFunc("/internal/translate/ids", handler.handlePostTranslateIDs).Methods("POST").Name("PostTranslateIDs") router.HandleFunc("/internal/index/{index}/field/{field}/attr/diff", handler.handlePostFieldAttrDiff).Methods("POST").Name("PostFieldAttrDiff") router.HandleFunc("/internal/index/{index}/field/{field}/remote-available-shards/{shardID}", handler.handleDeleteRemoteAvailableShard).Methods("DELETE") + router.HandleFunc("/internal/index/{index}/field/{field}/attr/data", handler.handleGetFieldAttrData).Methods("GET").Name("GetFieldAttrData") + router.HandleFunc("/internal/index/{index}/shard/{shard}/snapshot", handler.handleGetIndexShardSnapshot).Methods("GET").Name("GetIndexShardSnapshot") + router.HandleFunc("/internal/index/{index}/shards", handler.handleGetIndexAvailableShards).Methods("GET").Name("GetIndexAvailableShards") router.HandleFunc("/internal/nodes", handler.handleGetNodes).Methods("GET").Name("GetNodes") router.HandleFunc("/internal/shards/max", handler.handleGetShardsMax).Methods("GET").Name("GetShardsMax") // TODO: deprecate, but it's being used by the client @@ -443,6 +450,7 @@ func newRouter(handler *Handler) http.Handler { router.HandleFunc("/internal/idalloc/reserve", handler.handleReserveIDs).Methods("POST").Name("ReserveIDs") router.HandleFunc("/internal/idalloc/commit", handler.handleCommitIDs).Methods("POST").Name("CommitIDs") router.HandleFunc("/internal/idalloc/reset/{index}", handler.handleResetIDAlloc).Methods("POST").Name("ResetIDAlloc") + router.HandleFunc("/internal/idalloc/data", handler.handleIDAllocData).Methods("GET").Name("IDAllocData") // endpoints for collecting cpu profiles from a chosen begin point to // when the client wants to stop. Used for profiling imports that @@ -976,6 +984,30 @@ func (h *Handler) handleCPUProfileStop(w http.ResponseWriter, r *http.Request) { } } +// handleGetIndexAvailableShards handles GET /internal/index/:index/shards requests. +func (h *Handler) handleGetIndexAvailableShards(w http.ResponseWriter, r *http.Request) { + if !validHeaderAcceptJSON(r.Header) { + http.Error(w, "JSON only acceptable response", http.StatusNotAcceptable) + return + } + + indexName := mux.Vars(r)["index"] + shards, err := h.api.AvailableShards(r.Context(), indexName) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(getIndexAvailableShardsResponse{Shards: shards.Slice()}); err != nil { + h.logger.Errorf("write shards-max response error: %s", err) + } +} + +type getIndexAvailableShardsResponse struct { + Shards []uint64 `json:"shards"` +} + // handleGetShardsMax handles GET /internal/shards/max requests. func (h *Handler) handleGetShardsMax(w http.ResponseWriter, r *http.Request) { if !validHeaderAcceptJSON(r.Header) { @@ -1190,6 +1222,14 @@ func (h *Handler) handlePostIndexAttrDiff(w http.ResponseWriter, r *http.Request } } +// handleGetIndexAttrData handles GET /internal/index/{index}/attr/data requests. +func (h *Handler) handleGetIndexAttrData(w http.ResponseWriter, r *http.Request) { + if err := h.api.WriteColumnAttrDataTo(r.Context(), w, mux.Vars(r)["index"]); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } +} + func (h *Handler) handleGetActiveQueries(w http.ResponseWriter, r *http.Request) { var rtype string switch { @@ -1714,6 +1754,42 @@ type postFieldAttrDiffResponse struct { Attrs map[uint64]map[string]interface{} `json:"attrs"` } +// handleGetFieldAttrData handles GET /internal/index/{index}/field/{field}/attr/data requests. +func (h *Handler) handleGetFieldAttrData(w http.ResponseWriter, r *http.Request) { + if err := h.api.WriteRowAttrDataTo(r.Context(), w, mux.Vars(r)["index"], mux.Vars(r)["field"]); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } +} + +// handleGetIndexShardSnapshot handles GET /internal/index/{index}/shard/{shard}/snapshot requests. +func (h *Handler) handleGetIndexShardSnapshot(w http.ResponseWriter, r *http.Request) { + indexName := mux.Vars(r)["index"] + shard, err := strconv.ParseUint(mux.Vars(r)["shard"], 10, 64) + if err != nil { + http.Error(w, "Invalid shard parameter", http.StatusBadRequest) + return + } + + rc, err := h.api.IndexShardSnapshot(r.Context(), indexName, shard) + if err != nil { + switch errors.Cause(err) { + case pilosa.ErrIndexNotFound: + http.Error(w, err.Error(), http.StatusNotFound) + default: + http.Error(w, err.Error(), http.StatusInternalServerError) + } + return + } + defer rc.Close() + + // 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 + } +} + // readQueryRequest parses an query parameters from r. func (h *Handler) readQueryRequest(r *http.Request) (*pilosa.QueryRequest, error) { switch r.Header.Get("Content-Type") { @@ -1724,6 +1800,16 @@ func (h *Handler) readQueryRequest(r *http.Request) (*pilosa.QueryRequest, error } } +// 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) +} + // readProtobufQueryRequest parses query parameters in protobuf from r. func (h *Handler) readProtobufQueryRequest(r *http.Request) (*pilosa.QueryRequest, error) { // Slurp the body. @@ -2023,13 +2109,29 @@ func (h *Handler) handleGetFragmentData(w http.ResponseWriter, r *http.Request) // handleGetTranslateData handles GET /internal/translate/data requests. func (h *Handler) handleGetTranslateData(w http.ResponseWriter, r *http.Request) { - // Read partition parameter. q := r.URL.Query() + + // Perform field translation copy, if field specified. + if fieldName := q.Get("field"); fieldName != "" { + // Retrieve field data from holder. + p, err := h.api.FieldTranslateData(r.Context(), q.Get("index"), fieldName) + if err != nil { + http.Error(w, err.Error(), http.StatusNotFound) + return + } + // Stream translate data to response body. + if _, err := p.WriteTo(w); err != nil { + h.logger.Errorf("error streaming translation data: %s", err) + } + } + + // Otherwise read partition parameter for index translation copy. partition, err := strconv.ParseUint(q.Get("partition"), 10, 32) if err != nil { - http.Error(w, "partition required", http.StatusBadRequest) + http.Error(w, "partition or field required", http.StatusBadRequest) return } + // Retrieve partition data from holder. p, err := h.api.TranslateData(r.Context(), q.Get("index"), int(partition)) if err != nil { @@ -2990,3 +3092,11 @@ func (h *Handler) handleResetIDAlloc(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) w.Write([]byte("OK")) //nolint:errcheck } + +func (h *Handler) handleIDAllocData(w http.ResponseWriter, r *http.Request) { + w.Header().Add("Content-Type", "application/octet-stream") + if err := h.api.WriteIDAllocDataTo(w); err != nil { + http.Error(w, fmt.Sprintf("writeing id allocation data: %v", err.Error()), http.StatusInternalServerError) + return + } +} diff --git a/idalloc.go b/idalloc.go index 97657f6e4..c9c954d76 100644 --- a/idalloc.go +++ b/idalloc.go @@ -17,6 +17,7 @@ package pilosa import ( "encoding/binary" "fmt" + "io" "math/bits" "sort" "time" @@ -72,6 +73,20 @@ func (ida *idAllocator) Close() error { return ida.db.Close() } +func (ida *idAllocator) WriteTo(w io.Writer) (int64, error) { + if ida == nil || ida.db == nil { + return 0, fmt.Errorf("idAllocator closed") + } + + tx, err := ida.db.Begin(false) + if err != nil { + return 0, err + } + defer func() { _ = tx.Rollback() }() + + return tx.WriteTo(w) +} + // ErrIDOffsetDesync is an error generated when attempting to reserve IDs at a committed offset. // This will typically happen when kafka partitions are moved between kafka ingesters - there may be a brief period in which 2 ingesters are processing the same messages at the same time. // The ingester can resolve this by ignoring messages under base. diff --git a/pilosa_internal_test.go b/pilosa_internal_test.go index aca9c6e92..3de8cfbbf 100644 --- a/pilosa_internal_test.go +++ b/pilosa_internal_test.go @@ -16,6 +16,7 @@ package pilosa import ( "bytes" + "io" "reflect" "testing" @@ -73,6 +74,8 @@ func (s *memAttrStore) BlockData(i uint64) (map[uint64]map[string]interface{}, e return nil, nil } +func (s *memAttrStore) WriteTo(w io.Writer) (int64, error) { return 0, nil } + func TestAPI_CombineForExistence(t *testing.T) { bm := roaring.NewBitmap(pos(1, 1), pos(1, 2), pos(1, 3), pos(1, 65537), pos(1, 65538), pos(2, 1), pos(2, 2), pos(2, 5), pos(2, 65537), pos(2, 65538)) buf := new(bytes.Buffer) diff --git a/rbf.go b/rbf.go index fbbfb36e8..ff2eea25c 100644 --- a/rbf.go +++ b/rbf.go @@ -440,6 +440,11 @@ func (tx *RBFTx) GetFieldSizeBytes(index, field string) (uint64, error) { return tx.tx.GetSizeBytesWithPrefix(string(txkey.FieldPrefix(index, field))) } +// SnapshotReader returns a reader that provides a snapshot of the current database. +func (tx *RBFTx) SnapshotReader() (io.Reader, error) { + return tx.tx.SnapshotReader() +} + // rbfName returns a NULL-separated key used for identifying bitmap maps in RBF. func rbfName(index, field, view string, shard uint64) string { return string(txkey.Prefix(index, field, view, shard)) diff --git a/rbf/tx.go b/rbf/tx.go index 2996195bc..85b4af03c 100644 --- a/rbf/tx.go +++ b/rbf/tx.go @@ -2079,6 +2079,40 @@ func (tx *Tx) GetSortedFieldViewList() (fvs []txkey.FieldView, _ error) { return } +// SnapshotReader returns a reader that provides a snapshot for the current database state. +func (tx *Tx) SnapshotReader() (io.Reader, error) { + if tx.db == nil { + return nil, ErrTxClosed + } + return &snapshotReader{tx: tx}, nil +} + +type snapshotReader struct { + tx *Tx + pgno uint32 +} + +func (r *snapshotReader) Read(p []byte) (n int, err error) { + // Exit if we are past the end of the database. + if r.pgno >= readMetaPageN(r.tx.meta[:]) { + return 0, io.EOF + } + + // Otherwise look up the page data from mmap or page cache and copy it out. + buf, _, err := r.tx.readPage(r.pgno) + if err != nil { + return 0, err + } else if len(p) < len(buf) { + return 0, io.ErrShortBuffer + } + copy(p, buf) + + // Increment the page number. + r.pgno++ + + return len(buf), nil +} + type PageInfo interface { pageInfo() }