mirror of
https://github.com/hanzoai/traces.git
synced 2026-10-30 18:20:41 +00:00
Zach LeslieandGitHub
c5fe2e7e6b
* [backend-scheduler] initial work file sharding for more targeted marshaling * Add integration coverage for both work implementations * Update config manifest * Add integration coverage for sharded work configuration * Fix integration test config path * Add test coverage for sharded work migrations and marshaling * Use new sharded work file in the backend and handle migration * Clean up backend after migration * Use shared method for shard file name * Add method to write files atomically and test * Update work interface for simpler flush/load interactions * Update callers for local flush/load interface methods * Fix lint * Fix lint * Fix lint * Clean up * Clean up after refactor new interface functions * Promote the sharded implementation to be the only implementation * Clean up comments * Generate config manifest * Clean up interface and improve test coverage * Drop transitional comments * Protect work marshal/unmarshal with a mutex * Add test coverage for concurrent marshal/unmarshal * Lint * Fix race condition in Marshal/Unmarshal when jobs are operating * Relocate methods * Update changelog * Drop interface method only used in tests * Rename method for clarity, add doc * Update doc for clarity * Lock for safety * Clarify and improve locking flush behavior * Encapsulate the lock/unlock in a func * Address some PR feedback * Further unifi shard file name generation * Modify work to use int rather than uint8 * Fix lint
102 lines
2.6 KiB
Go
102 lines
2.6 KiB
Go
package backendscheduler
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
"path/filepath"
|
|
|
|
"github.com/go-kit/log/level"
|
|
"github.com/grafana/tempo/modules/backendscheduler/work"
|
|
"github.com/grafana/tempo/pkg/util/log"
|
|
"github.com/grafana/tempo/tempodb/backend"
|
|
)
|
|
|
|
func (s *BackendScheduler) flushWorkCacheToBackend(ctx context.Context) error {
|
|
_, span := tracer.Start(ctx, "flushWorkCacheToBackend")
|
|
defer span.End()
|
|
|
|
s.mtx.Lock()
|
|
defer s.mtx.Unlock()
|
|
|
|
b, err := s.work.Marshal()
|
|
if err != nil {
|
|
metricWorkFlushesFailed.Inc()
|
|
return fmt.Errorf("failed to marshal work cache: %w", err)
|
|
}
|
|
|
|
err = s.writer.Write(ctx, backend.WorkFileName, []string{}, bytes.NewReader(b), int64(len(b)), nil)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to flush work cache: %w", err)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// loadWorkCache loads the work cache
|
|
func (s *BackendScheduler) loadWorkCache(ctx context.Context) error {
|
|
ctx, span := tracer.Start(ctx, "loadWorkCache")
|
|
defer span.End()
|
|
|
|
// Check if shard files exist before attempting to load them
|
|
if s.checkForShardFiles() {
|
|
// Try to load the local work cache first, falling back to the backend if it doesn't exist.
|
|
err := s.work.LoadFromLocal(ctx, s.cfg.LocalWorkPath)
|
|
if err != nil {
|
|
if !os.IsNotExist(err) {
|
|
level.Error(log.Logger).Log("msg", "failed to read work cache from local path", "path", s.cfg.LocalWorkPath, "error", err)
|
|
}
|
|
|
|
return s.loadWorkCacheFromBackend(ctx)
|
|
}
|
|
|
|
return s.replayWorkOnBlocklist(ctx)
|
|
}
|
|
|
|
// No shard files found locally, load from backend
|
|
return s.loadWorkCacheFromBackend(ctx)
|
|
}
|
|
|
|
func (s *BackendScheduler) loadWorkCacheFromBackend(ctx context.Context) error {
|
|
ctx, span := tracer.Start(ctx, "loadWorkCacheFromBackend")
|
|
defer span.End()
|
|
|
|
reader, _, err := s.reader.Read(ctx, backend.WorkFileName, backend.KeyPath{}, nil)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer func() {
|
|
if closeErr := reader.Close(); err != nil {
|
|
level.Error(log.Logger).Log("msg", "failed to close reader", "err", closeErr)
|
|
}
|
|
}()
|
|
|
|
data, err := io.ReadAll(reader)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
err = s.work.Unmarshal(data)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to unmarshal work cache: %w", err)
|
|
}
|
|
|
|
return s.replayWorkOnBlocklist(ctx)
|
|
}
|
|
|
|
// checkForShardFiles checks if any shard files exist
|
|
func (s *BackendScheduler) checkForShardFiles() bool {
|
|
for i := range work.ShardCount {
|
|
if _, err := os.Stat(s.filepathForShard(uint8(i))); err == nil {
|
|
return true // Found at least one shard file
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func (s *BackendScheduler) filepathForShard(shardID uint8) string {
|
|
return filepath.Join(s.cfg.LocalWorkPath, work.FileNameForShard(shardID))
|
|
}
|