111 lines
3.5 KiB
Go
111 lines
3.5 KiB
Go
// Copyright 2023 PingCAP, 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 asyncloaddata
|
|
|
|
import (
|
|
"encoding/json"
|
|
"sync"
|
|
|
|
"go.uber.org/atomic"
|
|
"golang.org/x/exp/maps"
|
|
)
|
|
|
|
// LogicalImportProgress is the progress info of the logical import mode.
|
|
type LogicalImportProgress struct {
|
|
// LoadedFileSize is the size of the data that's loaded in bytes. It's
|
|
// larger than the actual loaded data size, but due to the fact that reading
|
|
// is once-a-block and a block may generate multiple tasks that are
|
|
// concurrently executed, we can't know the actual loaded data size easily.
|
|
LoadedFileSize atomic.Int64
|
|
}
|
|
|
|
// PhysicalImportProgress is the progress info of the physical import mode.
|
|
type PhysicalImportProgress struct {
|
|
// ReadRowCnt is the number of rows read from data files.
|
|
// Lines ignored by IGNORE N LINES clause is not included.
|
|
ReadRowCnt atomic.Uint64
|
|
// EncodeFileSize is the size of the file that has finished KV encoding in bytes.
|
|
// it should equal to SourceFileSize eventually.
|
|
EncodeFileSize atomic.Int64
|
|
|
|
colSizeMu sync.Mutex
|
|
colSizeMap map[int64]int64
|
|
}
|
|
|
|
// Progress is the progress of the LOAD DATA task.
|
|
type Progress struct {
|
|
// SourceFileSize is the size of the source file in bytes. When we can't get
|
|
// the size of the source file, it will be set to -1.
|
|
// Currently, the value is read by seek(0, end), when LOAD DATA LOCAL we wrap
|
|
// SimpleSeekerOnReadCloser on MySQL client connection which doesn't support
|
|
// it.
|
|
SourceFileSize int64
|
|
*LogicalImportProgress `json:",inline"`
|
|
*PhysicalImportProgress `json:",inline"`
|
|
// LoadedRowCnt is the number of rows that has been loaded.
|
|
// for physical mode, it's the number of rows that has been imported into TiKV.
|
|
// in SHOW LOAD JOB we call it Imported_Rows, to make it compatible with 7.0,
|
|
// the variable name is not changed.
|
|
LoadedRowCnt atomic.Uint64
|
|
}
|
|
|
|
// NewProgress creates a new Progress.
|
|
// todo: better pass import mode, but it causes import cycle.
|
|
func NewProgress(logicalImport bool) *Progress {
|
|
var li *LogicalImportProgress
|
|
var pi *PhysicalImportProgress
|
|
if logicalImport {
|
|
li = &LogicalImportProgress{}
|
|
} else {
|
|
pi = &PhysicalImportProgress{
|
|
colSizeMap: make(map[int64]int64),
|
|
}
|
|
}
|
|
return &Progress{
|
|
SourceFileSize: -1,
|
|
LogicalImportProgress: li,
|
|
PhysicalImportProgress: pi,
|
|
}
|
|
}
|
|
|
|
// AddColSize adds the size of the column to the progress.
|
|
func (p *Progress) AddColSize(colSizeMap map[int64]int64) {
|
|
p.colSizeMu.Lock()
|
|
defer p.colSizeMu.Unlock()
|
|
for key, value := range colSizeMap {
|
|
p.colSizeMap[key] += value
|
|
}
|
|
}
|
|
|
|
// GetColSize returns the size of the column.
|
|
func (p *Progress) GetColSize() map[int64]int64 {
|
|
p.colSizeMu.Lock()
|
|
defer p.colSizeMu.Unlock()
|
|
return maps.Clone(p.colSizeMap)
|
|
}
|
|
|
|
// String implements the fmt.Stringer interface.
|
|
func (p *Progress) String() string {
|
|
bs, _ := json.Marshal(p)
|
|
return string(bs)
|
|
}
|
|
|
|
// ProgressFromJSON creates Progress from a JSON string.
|
|
func ProgressFromJSON(bs []byte) (*Progress, error) {
|
|
var p Progress
|
|
err := json.Unmarshal(bs, &p)
|
|
return &p, err
|
|
}
|