Repository: pingcap/tidb
File: pkg/disttask/framework/metering/metering.go (Go)

Code before:
}

// Meter is responsible for recording and reporting metering data.
type Meter struct {
	sync.Mutex
	recorders map[int64]*wrappedRecorder
	// taskID -> last flushed data
	// when flushing, we scrape the latest data from recorders and calculate the
	// delta and write to the metering storage.
	lastFlushedData map[int64]*Data
	uuid            string
	writer          meteringwriterapi.MeteringWriter
	logger          *zap.Logger
}

// NewMeter creates a new Meter instance.
func NewMeter(cfg *mconfig.MeteringConfig) (*Meter, error) {
	logger := logutil.BgLogger().With(zap.String("component", "meter"))
	if len(cfg.Type) == 0 || len(cfg.Bucket) == 0 {
		return nil, nil
	}
	s3Config := cfg.ToProviderConfig()
	provider, err := storage.NewObjectStorageProvider(s3Config)
	if err != nil {
		return nil, errors.Wrap(err, "failed to create storage provider")
	}
	// if there are network issues, we might successfully write the metering data,
	// but the writer still returns error, we will retry write with the same TS
	// in this case, which means the metering data file will have the same name,
	// we set WithOverwriteExisting to true to avoid the retry write report error
	// in this case.
	// IgnoreExisting wound be more appropriate, but the SDK doesn't provide it.
	meteringConfig := mconfig.DefaultConfig().WithLogger(logger).WithOverwriteExisting(true)
	writer := meteringwriter.NewMeteringWriterFromConfig(provider, meteringConfig, cfg)
	return newMeterWithWriter(logger, writer), nil
}

func newMeterWithWriter(logger *zap.Logger, writer meteringwriterapi.MeteringWriter) *Meter {
	return &Meter{
		logger:          logger,
		recorders:       make(map[int64]*wrappedRecorder),
		lastFlushedData: make(map[int64]*Data),
		writer:          writer,
		uuid:            strings.ReplaceAll(uuid.New().String(), "-", "_"), // no dash in the metering sdk
	}
}

func (m *Meter) getOrRegisterRecorder(r *Recorder) *Recorder {
	m.Lock()
	defer m.Unlock()
	if old, ok := m.recorders[r.taskID]; ok {

Review comment (bug):
The unconditional overwrite at line 120 could reset the retry count if `addFailedData` is called with a timestamp that already exists in the map. While the current implementation uses unique timestamps per flush interval, a safer approach would be to check if the timestamp already exists and skip adding it again, or log a warning. This prevents accidental retry count resets if the timing logic changes in the future.
```suggestion
	}
	if _, exists := d.data[ts]; exists {
		logutil.BgLogger().Warn(
			"retryData.addFailedData called with existing timestamp; skipping to avoid resetting retry count",
			zap.Int64("ts", ts),
		)
		return
	}
```

Code after:
	meteringInstance.Store(m)
}

type wrappedRecorder struct {
	*Recorder
	unregistered bool
}

type writeFailData struct {
	ts       int64
	retryCnt int
	items    []map[string]any
}

type retryData struct {
	mu sync.Mutex
	// TS -> write failed data
	data map[int64]*writeFailData
}

func (d *retryData) addFailedData(ts int64, items []map[string]any) {
	d.mu.Lock()
	defer d.mu.Unlock()
	if d.data == nil {
		d.data = make(map[int64]*writeFailData)
	}
	d.data[ts] = &writeFailData{ts: ts, items: items}
}

func (d *retryData) getDataClone() map[int64]*writeFailData {
	d.mu.Lock()
	defer d.mu.Unlock()
	return maps.Clone(d.data)
}

func (d *retryData) remove(needRemove []*writeFailData) {
	if len(needRemove) == 0 {
		return
	}
	d.mu.Lock()
	defer d.mu.Unlock()
	for _, wd := range needRemove {
		delete(d.data, wd.ts)
	}
}

// Meter is responsible for recording and reporting metering data.
type Meter struct {
	mu        sync.Mutex
	recorders map[int64]*wrappedRecorder
	// taskID -> last flushed data