wehub-resource-sync
469 lines
13 KiB
Go
469 lines
13 KiB
Go
// Copyright 2019 Dolthub, 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 pull
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"log"
|
|
"math"
|
|
"os"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"golang.org/x/sync/errgroup"
|
|
|
|
"github.com/dolthub/dolt/go/libraries/doltcore/dconfig"
|
|
"github.com/dolthub/dolt/go/store/chunks"
|
|
"github.com/dolthub/dolt/go/store/hash"
|
|
"github.com/dolthub/dolt/go/store/nbs"
|
|
)
|
|
|
|
// ErrDBUpToDate is the error code returned from NewPuller in the event that there is no work to do.
|
|
var ErrDBUpToDate = errors.New("the database does not need to be pulled as it's already up to date")
|
|
|
|
// ErrIncompatibleSourceChunkStore is the error code returned from NewPuller in
|
|
// the event that the source ChunkStore does not implement `NBSCompressedChunkStore`.
|
|
var ErrIncompatibleSourceChunkStore = errors.New("the chunk store of the source database does not implement NBSCompressedChunkStore.")
|
|
|
|
type WalkAddrs func(chunks.Chunk, func(hash.Hash, bool) error) error
|
|
|
|
// Puller is used to sync data between to Databases
|
|
type Puller struct {
|
|
waf WalkAddrs
|
|
|
|
srcChunkStore nbs.NBSCompressedChunkStore
|
|
sinkDBCS chunks.ChunkStore
|
|
hashes hash.HashSet
|
|
|
|
wr *PullTableFileWriter
|
|
|
|
tempDir string
|
|
pushLog *log.Logger
|
|
|
|
statsCh chan Stats
|
|
stats *stats
|
|
}
|
|
|
|
// NewPuller creates a new Puller instance to do the syncing. If a nil puller is returned without error that means
|
|
// that there is nothing to pull and the sinkDB is already up to date.
|
|
func NewPuller(
|
|
ctx context.Context,
|
|
tempDir string,
|
|
targetFileSz uint64,
|
|
srcCS, sinkCS chunks.ChunkStore,
|
|
walkAddrs WalkAddrs,
|
|
hashes []hash.Hash,
|
|
statsCh chan Stats,
|
|
) (*Puller, error) {
|
|
// Sanity Check
|
|
hs := hash.NewHashSet(hashes...)
|
|
missing, err := srcCS.HasMany(ctx, hs)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if missing.Size() != 0 {
|
|
return nil, errors.New("not found")
|
|
}
|
|
|
|
hs = hash.NewHashSet(hashes...)
|
|
missing, err = sinkCS.HasMany(ctx, hs)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if missing.Size() == 0 {
|
|
return nil, ErrDBUpToDate
|
|
}
|
|
|
|
if srcCS.Version() != sinkCS.Version() {
|
|
return nil, fmt.Errorf("cannot pull from src to sink; src version is %v and sink version is %v", srcCS.Version(), sinkCS.Version())
|
|
}
|
|
|
|
srcChunkStore, ok := srcCS.(nbs.NBSCompressedChunkStore)
|
|
if !ok {
|
|
return nil, ErrIncompatibleSourceChunkStore
|
|
}
|
|
|
|
// walkAddrs can be used for getAddrs on the AddTableFile call as well.
|
|
getAddrs := func(c chunks.Chunk) chunks.InsertAddrsCb {
|
|
return func(ctx context.Context, ins hash.HashSet, filter chunks.PendingRefExists) error {
|
|
return walkAddrs(c, func(h hash.Hash, _ bool) error {
|
|
if !filter(h) {
|
|
ins.Insert(h)
|
|
}
|
|
return nil
|
|
})
|
|
}
|
|
}
|
|
|
|
wr := NewPullTableFileWriter(PullTableFileWriterConfig{
|
|
ConcurrentUploads: 2,
|
|
TargetFileSize: targetFileSz,
|
|
MaximumBufferedFiles: 8,
|
|
TempDir: tempDir,
|
|
DestStore: sinkCS.(chunks.TableFileStore),
|
|
GetAddrs: getAddrs,
|
|
})
|
|
|
|
p := &Puller{
|
|
waf: walkAddrs,
|
|
srcChunkStore: srcChunkStore,
|
|
sinkDBCS: sinkCS,
|
|
hashes: hash.NewHashSet(hashes...),
|
|
wr: wr,
|
|
tempDir: tempDir,
|
|
statsCh: statsCh,
|
|
stats: &stats{
|
|
wrStatsGetter: wr.GetStats,
|
|
fetcher: &fetchStatsRecorder{},
|
|
},
|
|
}
|
|
|
|
// Route chunk store debug logging into the push log. For a push, the sink
|
|
// is the remote store; for a fetch, the source is. Wire both so either
|
|
// operation's remote-side messages are captured.
|
|
if lcs, ok := sinkCS.(chunks.LoggingChunkStore); ok {
|
|
lcs.SetLogger(p)
|
|
}
|
|
if lcs, ok := srcCS.(chunks.LoggingChunkStore); ok {
|
|
lcs.SetLogger(p)
|
|
}
|
|
|
|
return p, nil
|
|
}
|
|
|
|
func (p *Puller) Logf(fmt string, args ...interface{}) {
|
|
if p.pushLog != nil {
|
|
p.pushLog.Printf(fmt, args...)
|
|
}
|
|
}
|
|
|
|
type readable interface {
|
|
Reader() (io.ReadCloser, error)
|
|
Remove() error
|
|
}
|
|
|
|
type tempTblFile struct {
|
|
read readable
|
|
id string
|
|
contentHash []byte
|
|
numChunks int
|
|
chunksLen uint64
|
|
contentLen uint64
|
|
pending io.Closer
|
|
}
|
|
|
|
type countingReader struct {
|
|
io.ReadCloser
|
|
cnt *uint64
|
|
}
|
|
|
|
func (c countingReader) Read(p []byte) (int, error) {
|
|
n, err := c.ReadCloser.Read(p)
|
|
atomic.AddUint64(c.cnt, uint64(n))
|
|
return n, err
|
|
}
|
|
|
|
func emitStats(s *stats, ch chan Stats, logf func(string, ...interface{})) (cancel func()) {
|
|
done := make(chan struct{})
|
|
var wg sync.WaitGroup
|
|
wg.Add(2)
|
|
cancel = func() {
|
|
close(done)
|
|
wg.Wait()
|
|
}
|
|
|
|
go func() {
|
|
defer wg.Done()
|
|
sampleduration := 100 * time.Millisecond
|
|
samplesinsec := uint64((1 * time.Second) / sampleduration)
|
|
weight := 0.1
|
|
ticker := time.NewTicker(sampleduration)
|
|
defer ticker.Stop()
|
|
var lastSendBytes, lastFetchedBytes uint64
|
|
for {
|
|
select {
|
|
case <-ticker.C:
|
|
wrStats := s.wrStatsGetter()
|
|
newSendBytes := wrStats.FinishedSendBytes
|
|
newFetchedBytes := atomic.LoadUint64(&s.fetchedSourceBytes)
|
|
sendBytesDiff := newSendBytes - lastSendBytes
|
|
fetchedBytesDiff := newFetchedBytes - lastFetchedBytes
|
|
|
|
newSendBPS := float64(sendBytesDiff * samplesinsec)
|
|
newFetchedBPS := float64(fetchedBytesDiff * samplesinsec)
|
|
|
|
curSendBPS := math.Float64frombits(atomic.LoadUint64(&s.sendBytesPerSec))
|
|
curFetchedBPS := math.Float64frombits(atomic.LoadUint64(&s.fetchedSourceBytesPerSec))
|
|
|
|
smoothedSendBPS := newSendBPS
|
|
if curSendBPS != 0 {
|
|
smoothedSendBPS = curSendBPS + weight*(newSendBPS-curSendBPS)
|
|
}
|
|
|
|
smoothedFetchBPS := newFetchedBPS
|
|
if curFetchedBPS != 0 {
|
|
smoothedFetchBPS = curFetchedBPS + weight*(newFetchedBPS-curFetchedBPS)
|
|
}
|
|
|
|
if smoothedSendBPS < 1 {
|
|
smoothedSendBPS = 0
|
|
}
|
|
if smoothedFetchBPS < 1 {
|
|
smoothedFetchBPS = 0
|
|
}
|
|
|
|
atomic.StoreUint64(&s.sendBytesPerSec, math.Float64bits(smoothedSendBPS))
|
|
atomic.StoreUint64(&s.fetchedSourceBytesPerSec, math.Float64bits(smoothedFetchBPS))
|
|
|
|
lastSendBytes = newSendBytes
|
|
lastFetchedBytes = newFetchedBytes
|
|
case <-done:
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
|
|
go func() {
|
|
defer wg.Done()
|
|
updateduration := 1 * time.Second
|
|
ticker := time.NewTicker(updateduration)
|
|
defer ticker.Stop()
|
|
emit := func(st Stats) {
|
|
if ch != nil {
|
|
ch <- st
|
|
}
|
|
fs := s.fetcher.read()
|
|
compBytes := atomic.LoadUint64(&s.fetchedSourceCompBytes)
|
|
// Source side counters:
|
|
// chunks - wanted chunks received / total requested
|
|
// uncompressed - decompressed size of the wanted chunks
|
|
// compressed - on-wire (snappy) size of just the wanted chunks
|
|
// wire - bytes actually downloaded, including dark bytes
|
|
// pulled in by range coalescing (wire - compressed
|
|
// ~= coalescing overhead)
|
|
// fetch_rate - uncompressed bytes/sec
|
|
// range_downloads - completed coalesced range downloads (and retries)
|
|
// inflight/peak - concurrent range downloads in flight
|
|
// Sink side counters:
|
|
// sent/buffered - table file bytes written to / queued for the dest
|
|
// send_rate - sent bytes/sec
|
|
logf("stats: source[chunks=%d/%d uncompressed=%dB compressed=%dB wire=%dB fetch_rate=%.0fB/s range_downloads=%d retries=%d inflight=%d peak=%d] sink[sent=%dB buffered=%dB send_rate=%.0fB/s]",
|
|
st.FetchedSourceChunks, st.TotalSourceChunks,
|
|
st.FetchedSourceBytes, compBytes, fs.DownloadedBytes, st.FetchedSourceBytesPerSec,
|
|
fs.CompletedDownloads, fs.Retries, fs.InFlight, fs.PeakInFlight,
|
|
st.FinishedSendBytes, st.BufferedSendBytes, st.SendBytesPerSec)
|
|
}
|
|
for {
|
|
select {
|
|
case <-ticker.C:
|
|
emit(s.read())
|
|
case <-done:
|
|
emit(s.read())
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
|
|
return cancel
|
|
}
|
|
|
|
type stats struct {
|
|
wrStatsGetter func() PullTableFileWriterStats
|
|
sendBytesPerSec uint64
|
|
totalSourceChunks uint64
|
|
fetchedSourceChunks uint64
|
|
fetchedSourceBytes uint64
|
|
fetchedSourceCompBytes uint64
|
|
fetchedSourceBytesPerSec uint64
|
|
sendBytesPerSecF float64
|
|
fetchedSourceBytesPerSecF float64
|
|
|
|
// fetcher accumulates download stats reported by the source ChunkFetcher
|
|
// (bytes fetched over the wire including dark bytes, download concurrency).
|
|
fetcher *fetchStatsRecorder
|
|
}
|
|
|
|
type Stats struct {
|
|
FinishedSendBytes uint64
|
|
BufferedSendBytes uint64
|
|
SendBytesPerSec float64
|
|
|
|
TotalSourceChunks uint64
|
|
FetchedSourceChunks uint64
|
|
FetchedSourceBytes uint64
|
|
FetchedSourceBytesPerSec float64
|
|
}
|
|
|
|
func (s *stats) read() Stats {
|
|
wrStats := s.wrStatsGetter()
|
|
|
|
var ret Stats
|
|
ret.FinishedSendBytes = wrStats.FinishedSendBytes
|
|
ret.BufferedSendBytes = wrStats.BufferedSendBytes
|
|
ret.SendBytesPerSec = math.Float64frombits(atomic.LoadUint64(&s.sendBytesPerSec))
|
|
ret.TotalSourceChunks = atomic.LoadUint64(&s.totalSourceChunks)
|
|
ret.FetchedSourceChunks = atomic.LoadUint64(&s.fetchedSourceChunks)
|
|
ret.FetchedSourceBytes = atomic.LoadUint64(&s.fetchedSourceBytes)
|
|
ret.FetchedSourceBytesPerSec = math.Float64frombits(atomic.LoadUint64(&s.fetchedSourceBytesPerSec))
|
|
return ret
|
|
}
|
|
|
|
func WithDiscardingStatsCh(cb func(statsCh chan Stats)) {
|
|
statsCh := make(chan Stats)
|
|
var wg sync.WaitGroup
|
|
wg.Go(func() {
|
|
for range statsCh {
|
|
}
|
|
})
|
|
wg.Go(func() {
|
|
defer close(statsCh)
|
|
cb(statsCh)
|
|
})
|
|
wg.Wait()
|
|
}
|
|
|
|
// Pull executes the sync operation
|
|
func (p *Puller) Pull(ctx context.Context) error {
|
|
// If push logging is enabled, create a unique log file for this pull
|
|
// operation and close it when the pull completes.
|
|
if dbg, ok := os.LookupEnv(dconfig.EnvPushLog); ok && strings.EqualFold(dbg, "true") {
|
|
f, err := os.CreateTemp(p.tempDir, "push-*.log")
|
|
if err == nil {
|
|
p.pushLog = log.New(f, "", log.LstdFlags|log.Lmicroseconds)
|
|
p.Logf("starting pull of %d hashes", p.hashes.Size())
|
|
defer func() {
|
|
p.Logf("pull complete")
|
|
f.Close()
|
|
}()
|
|
}
|
|
}
|
|
|
|
// Run stats emission if anyone is consuming the stats channel or if we are
|
|
// logging stats to the push log.
|
|
if p.statsCh != nil || p.pushLog != nil {
|
|
c := emitStats(p.stats, p.statsCh, p.Logf)
|
|
defer c()
|
|
}
|
|
|
|
eg, ctx := errgroup.WithContext(ctx)
|
|
|
|
rd := GetChunkFetcher(ctx, p.srcChunkStore, p.stats.fetcher)
|
|
|
|
const batchSize = 64 * 1024
|
|
tracker := NewPullChunkTracker(TrackerConfig{
|
|
BatchSize: batchSize,
|
|
HasManyer: p.sinkDBCS,
|
|
})
|
|
|
|
eg.Go(func() error {
|
|
return tracker.Run(ctx, p.hashes)
|
|
})
|
|
|
|
eg.Go(func() error {
|
|
return p.wr.Run(ctx)
|
|
})
|
|
|
|
// One thread calls ChunkFetcher.Get on each batch.
|
|
eg.Go(func() error {
|
|
for {
|
|
toFetch, hasMore, err := tracker.GetChunksToFetch(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if !hasMore {
|
|
return rd.CloseSend()
|
|
}
|
|
|
|
atomic.AddUint64(&p.stats.totalSourceChunks, uint64(len(toFetch)))
|
|
err = rd.Get(ctx, toFetch)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
})
|
|
|
|
// One thread reads the received chunks, walks their addresses and writes them to the table file.
|
|
eg.Go(func() (err error) {
|
|
// Closing the tracker belongs here, because this is
|
|
// the goroutine that writes to the tracker state with
|
|
// |Seen| and |TickProcessed| calls. Once this thread
|
|
// finishes, there will be no further writes to the
|
|
// tracker state.
|
|
defer tracker.Close()
|
|
defer func() {
|
|
cerr := rd.Close()
|
|
err = errors.Join(err, cerr)
|
|
}()
|
|
for {
|
|
cChk, err := rd.Recv(ctx)
|
|
if err == io.EOF {
|
|
// This means the requesting thread
|
|
// successfully saw all chunk addresses and
|
|
// called CloseSend and that all requested
|
|
// chunks were successfully delivered to this
|
|
// thread. Calling wr.Close() here will block
|
|
// on uploading any table files and will write
|
|
// the new table files to the destination's
|
|
// manifest.
|
|
p.wr.Close()
|
|
return nil
|
|
}
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if cChk.IsGhost() {
|
|
return fmt.Errorf("attempted to push or pull ghost chunk: %w", nbs.ErrGhostChunkRequested)
|
|
}
|
|
if cChk.IsEmpty() {
|
|
return errors.New("failed to get all chunks.")
|
|
}
|
|
|
|
atomic.AddUint64(&p.stats.fetchedSourceChunks, uint64(1))
|
|
// CompressedSize is the on-the-wire (snappy-compressed) size of
|
|
// just this wanted chunk, excluding any dark bytes pulled in by
|
|
// range coalescing (those are tracked by the fetcher recorder).
|
|
atomic.AddUint64(&p.stats.fetchedSourceCompBytes, uint64(cChk.CompressedSize()))
|
|
|
|
chnk, err := cChk.ToChunk()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
atomic.AddUint64(&p.stats.fetchedSourceBytes, uint64(len(chnk.Data())))
|
|
|
|
err = p.waf(chnk, func(h hash.Hash, _ bool) error {
|
|
tracker.Seen(ctx, h)
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
tracker.TickProcessed(ctx)
|
|
|
|
err = p.wr.AddToChunker(ctx, cChk)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
})
|
|
|
|
return eg.Wait()
|
|
}
|