github.com/dolthub/dolt/go@v0.40.5-0.20240520175717-68db7794bea6/libraries/utils/iohelp/read_with_stats.go (about)

     1  // Copyright 2021 Dolthub, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package iohelp
    16  
    17  import (
    18  	"io"
    19  	"sync/atomic"
    20  	"time"
    21  )
    22  
    23  const updateFrequency = 500 * time.Millisecond
    24  
    25  type ReadStats struct {
    26  	Read    uint64
    27  	Elapsed time.Duration
    28  	Percent float64
    29  }
    30  
    31  type ReaderWithStats struct {
    32  	read    uint64
    33  	size    int64
    34  	rd      io.Reader
    35  	start   time.Time
    36  	closeCh chan struct{}
    37  }
    38  
    39  func NewReaderWithStats(rd io.Reader, size int64) *ReaderWithStats {
    40  	return &ReaderWithStats{
    41  		rd:      rd,
    42  		size:    size,
    43  		closeCh: make(chan struct{}),
    44  	}
    45  }
    46  
    47  func (rws *ReaderWithStats) Start(updateFunc func(ReadStats)) {
    48  	rws.start = time.Now()
    49  	go func() {
    50  		timer := time.NewTimer(updateFrequency)
    51  		for {
    52  			select {
    53  			case <-rws.closeCh:
    54  				return
    55  			case <-timer.C:
    56  				read := atomic.LoadUint64(&rws.read)
    57  				elapsed := time.Since(rws.start)
    58  				var percent float64
    59  				if rws.size != 0 {
    60  					percent = float64(read) / float64(rws.size)
    61  				}
    62  				updateFunc(ReadStats{Read: read, Elapsed: elapsed, Percent: percent})
    63  				timer.Reset(updateFrequency)
    64  			}
    65  		}
    66  	}()
    67  }
    68  
    69  func (rws *ReaderWithStats) Close() error {
    70  	close(rws.closeCh)
    71  
    72  	if closer, ok := rws.rd.(io.Closer); ok {
    73  		return closer.Close()
    74  	}
    75  
    76  	return nil
    77  }
    78  
    79  func (rws *ReaderWithStats) Read(p []byte) (int, error) {
    80  	n, err := rws.rd.Read(p)
    81  
    82  	atomic.AddUint64(&rws.read, uint64(n))
    83  
    84  	return n, err
    85  }
    86  
    87  func (rws *ReaderWithStats) Size() int64 {
    88  	return rws.size
    89  }