mrproliu commented on code in PR #1145: URL: https://github.com/apache/skywalking-banyandb/pull/1145#discussion_r3329553954
########## banyand/backup/lifecycle/row_replay_measure.go: ########## @@ -0,0 +1,336 @@ +// Licensed to Apache Software Foundation (ASF) under one or more contributor +// license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright +// ownership. Apache Software Foundation (ASF) licenses this file to you 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 lifecycle + +import ( + "context" + "fmt" + "path/filepath" + "strconv" + "sync" + "sync/atomic" + "time" + + "google.golang.org/protobuf/types/known/timestamppb" + + "github.com/apache/skywalking-banyandb/api/common" + "github.com/apache/skywalking-banyandb/api/data" + commonv1 "github.com/apache/skywalking-banyandb/api/proto/banyandb/common/v1" + databasev1 "github.com/apache/skywalking-banyandb/api/proto/banyandb/database/v1" + measurev1 "github.com/apache/skywalking-banyandb/api/proto/banyandb/measure/v1" + "github.com/apache/skywalking-banyandb/banyand/internal/dump" + dumpmeasure "github.com/apache/skywalking-banyandb/banyand/internal/dump/measure" + "github.com/apache/skywalking-banyandb/banyand/metadata" + "github.com/apache/skywalking-banyandb/banyand/metadata/schema" + "github.com/apache/skywalking-banyandb/banyand/queue" + "github.com/apache/skywalking-banyandb/pkg/bus" + "github.com/apache/skywalking-banyandb/pkg/fs" + "github.com/apache/skywalking-banyandb/pkg/logger" + "github.com/apache/skywalking-banyandb/pkg/node" + "github.com/apache/skywalking-banyandb/pkg/partition" +) + +const ( + measureReplayBatchSize = 2000 + measureReplayBatchTimeout = 30 * time.Second +) + +// cachedMeasureSchema bundles a measure's schema with its derived entity and +// sharding-key routers so each row only requires one schema lookup per +// measure. shardingKey is non-nil only when the measure declares a non-empty +// ShardingKey; otherwise the entity-based shardID wins. +type cachedMeasureSchema struct { + schema *databasev1.Measure + shardingKey partition.Router + entityLocator partition.Locator +} + +// measureRowReplayer replays a group's measure parts row-by-row through the +// real Write API, resolving each row's schema on demand via EntityValues. +type measureRowReplayer struct { + selector node.Selector + client queue.Client + publisher queue.BatchPublisher + fs fs.FileSystem + logger *logger.Logger + measureSchemas map[string]*databasev1.Measure + schemaCache map[string]*cachedMeasureSchema + irCache map[string]*dump.IndexResolver + mergedRuleToTag map[uint32]dump.IndexedTagSpec + counter *uint64 + group string + batch []bus.Message + schemaCacheMu sync.Mutex + irCacheMu sync.Mutex + batchMu sync.Mutex + targetShardNum uint32 +} + +// newMeasureRowReplayer constructs a measure replayer, eagerly enumerating the +// group's measures and index rules for IndexResolver decoding. +func newMeasureRowReplayer( + ctx context.Context, + group string, targetShardNum uint32, + selector node.Selector, client queue.Client, + md metadata.Repo, fileSystem fs.FileSystem, + l *logger.Logger, counter *uint64, +) (*measureRowReplayer, error) { + measures, err := md.MeasureRegistry().ListMeasure(ctx, schema.ListOpt{Group: group}) + if err != nil { + return nil, fmt.Errorf("list measures of group %s: %w", group, err) + } + rules, err := md.IndexRuleRegistry().ListIndexRule(ctx, schema.ListOpt{Group: group}) + if err != nil { + return nil, fmt.Errorf("list index rules of group %s: %w", group, err) + } + bindings, err := md.IndexRuleBindingRegistry().ListIndexRuleBinding(ctx, schema.ListOpt{Group: group}) + if err != nil { + return nil, fmt.Errorf("list index rule bindings of group %s: %w", group, err) + } + // Index the already-listed measures by name so loadSchema resolves each + // subject from this snapshot instead of issuing a redundant GetMeasure + // per measure. Using one consistent snapshot also avoids mid-migration + // schema drift. + measureSchemas := make(map[string]*databasev1.Measure, len(measures)) + for _, m := range measures { + measureSchemas[m.GetMetadata().GetName()] = m + } + return &measureRowReplayer{ + group: group, + targetShardNum: targetShardNum, + selector: selector, + client: client, + publisher: client.NewBatchPublisher(measureReplayBatchTimeout), + fs: fileSystem, + logger: l, + measureSchemas: measureSchemas, + schemaCache: make(map[string]*cachedMeasureSchema), + irCache: make(map[string]*dump.IndexResolver), + mergedRuleToTag: deriveMergedRuleToTag(measures, rules, bindings), + counter: counter, + batch: make([]bus.Message, 0, measureReplayBatchSize), + }, nil +} + +// Close flushes pending messages, closes the publisher and releases cached +// IndexResolver handles. +func (r *measureRowReplayer) Close() (map[string]*common.Error, error) { + flushCtx, cancel := context.WithTimeout(context.Background(), measureReplayBatchTimeout) + defer cancel() + flushErr := r.flushBatch(flushCtx) + r.irCacheMu.Lock() + for _, ir := range r.irCache { + _ = ir.Close() + } + r.irCache = nil + r.irCacheMu.Unlock() + cee, closeErr := r.publisher.Close() + if flushErr != nil { + return cee, flushErr + } + return cee, closeErr +} + +// flushAndConfirm flushes the buffered batch, closes the publisher to obtain the +// per-node delivery result for everything sent since the last call, then opens a +// fresh publisher for subsequent parts. The batch publisher is client-streaming, +// so per-node errors are only observable once its stream closes; this lets a +// row-replayed part be confirmed durable before it is marked completed. +func (r *measureRowReplayer) flushAndConfirm(ctx context.Context) (map[string]*common.Error, error) { + flushErr := r.flushBatch(ctx) + cee, closeErr := r.publisher.Close() + r.publisher = r.client.NewBatchPublisher(measureReplayBatchTimeout) + if flushErr != nil { + return cee, flushErr + } + return cee, closeErr +} + +// loadSchema resolves a measure schema by name from the snapshot listed at +// construction and lazily caches its derived routers. +func (r *measureRowReplayer) loadSchema(measureName string) (*cachedMeasureSchema, error) { + r.schemaCacheMu.Lock() + defer r.schemaCacheMu.Unlock() + if c, ok := r.schemaCache[measureName]; ok { + return c, nil + } + m, ok := r.measureSchemas[measureName] + if !ok { + return nil, fmt.Errorf("measure schema %s/%s not found in group snapshot", r.group, measureName) + } + c := &cachedMeasureSchema{ + schema: m, + entityLocator: partition.NewEntityLocator(m.TagFamilies, m.Entity, m.GetMetadata().GetModRevision()), + } + if sk := m.GetShardingKey(); sk != nil && len(sk.GetTagNames()) > 0 { + c.shardingKey = partition.NewShardingKeyLocator(m.TagFamilies, sk) + } + r.schemaCache[measureName] = c + return c, nil +} + +// loadIndexResolver lazily opens an IndexResolver per source segment. +func (r *measureRowReplayer) loadIndexResolver(segmentPath string) (*dump.IndexResolver, error) { + r.irCacheMu.Lock() + defer r.irCacheMu.Unlock() + if ir, ok := r.irCache[segmentPath]; ok { + return ir, nil + } + ir, err := dump.NewIndexResolver(segmentPath, dump.DefaultIndexCacheSize, r.mergedRuleToTag) + if err != nil { + return nil, fmt.Errorf("open index resolver for %s: %w", segmentPath, err) + } + r.irCache[segmentPath] = ir Review Comment: Since all shards should walk from the same segment, I delete the map reference to a single reference to reduce the memory usage. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
