Skip to content

Commit

Permalink
enhance progress bar
Browse files Browse the repository at this point in the history
  • Loading branch information
@priyanshikhetwani
priyanshikhetwani committed Nov 11, 2024
1 parent 8cc8376 commit 410c6d9
Show file tree
Hide file tree
Showing 3 changed files with 96 additions and 1 deletion.
3 changes: 3 additions & 0 deletions go.mod
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,13 +20,15 @@ require (
github.com/spf13/cobra v1.8.1
github.com/spf13/pflag v1.0.5
github.com/stretchr/testify v1.9.0
github.com/vbauerster/mpb v3.4.0+incompatible
go.uber.org/mock v0.4.0
gopkg.in/yaml.v2 v2.4.0
k8s.io/klog/v2 v2.130.1
k8s.io/utils v0.0.0-20230726121419-3b25d923346b
)

require (
github.com/VividCortex/ewma v1.2.0 // indirect
github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2 // indirect
github.com/atotto/clipboard v0.1.4 // indirect
github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
Expand Down Expand Up @@ -88,6 +90,7 @@ require (
golang.org/x/net v0.28.0 // indirect
golang.org/x/sync v0.8.0 // indirect
golang.org/x/sys v0.25.0 // indirect
golang.org/x/term v0.23.0 // indirect
golang.org/x/text v0.18.0 // indirect
gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 // indirect
gopkg.in/yaml.v3 v3.0.1 // indirect
Expand Down
6 changes: 6 additions & 0 deletions go.sum
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,8 @@ github.com/MakeNowJust/heredoc v1.0.0/go.mod h1:mG5amYoWBHf8vpLOuehzbGGw0EHxpZZ6
github.com/OneOfOne/xxhash v1.2.2/go.mod h1:HSdplMjZKSmBqAxg5vPj2TmRDmfkzw+cTzAElWljhcU=
github.com/OpenPeeDeeP/depguard v1.0.1/go.mod h1:xsIw86fROiiwelg+jB2uM9PiKihMMmUx/1V+TNhjQvM=
github.com/StackExchange/wmi v0.0.0-20180116203802-5d049714c4a6/go.mod h1:3eOhrUMpNV+6aFIbp5/iudMxNCF27Vw2OZgy4xEx0Fg=
github.com/VividCortex/ewma v1.2.0 h1:f58SaIzcDXrSy3kWaHNvuJgJ3Nmz59Zji6XoJR/q1ow=
github.com/VividCortex/ewma v1.2.0/go.mod h1:nz4BbCtbLyFDeC9SUHbtcT5644juEuWfUAUnGx7j5l4=
github.com/alecthomas/template v0.0.0-20160405071501-a0175ee3bccc/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc=
github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
github.com/apparentlymart/go-cidr v1.1.0 h1:2mAhrMoF+nhXqxTzSZMUzDHkLjmIHC+Zzn4tdgBZjnU=
Expand Down Expand Up @@ -351,6 +353,8 @@ github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyC
github.com/valyala/fasthttp v1.2.0/go.mod h1:4vX61m6KN+xDduDNwXrhIAVZaZaZiQ1luJk8LWSxF3s=
github.com/valyala/quicktemplate v1.2.0/go.mod h1:EH+4AkTd43SvgIbQHYu59/cJyxDoOVRUAfrukLPuGJ4=
github.com/valyala/tcplisten v0.0.0-20161114210144-ceec8f93295a/go.mod h1:v3UYOV9WzVtRmSR+PDvWpU/qWl4Wa5LApYYX4ZtKbio=
github.com/vbauerster/mpb v3.4.0+incompatible h1:mfiiYw87ARaeRW6x5gWwYRUawxaW1tLAD8IceomUCNw=
github.com/vbauerster/mpb v3.4.0+incompatible/go.mod h1:zAHG26FUhVKETRu+MWqYXcI70POlC6N8up9p1dID7SU=
github.com/xiang90/probing v0.0.0-20190116061207-43a291ad63a2/go.mod h1:UETIi67q53MR2AWcXfiuqkDkRtnGDLqkBTpCHuJHxtU=
github.com/xordataexchange/crypt v0.0.3-0.20170626215501-b2862e3d0a77/go.mod h1:aYKd//L2LvnjZzWKhF00oedf4jCCReLcmhLdhm1A27Q=
github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
Expand Down Expand Up @@ -427,6 +431,8 @@ golang.org/x/sys v0.0.0-20210809222454-d867a43fc93e/go.mod h1:oPkhp1MJrh7nUepCBc
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.25.0 h1:r+8e+loiHxRqhXVl6ML1nO3l1+oFoWbnlu2Ehimmi34=
golang.org/x/sys v0.25.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/term v0.23.0 h1:F6D4vR+EHoL9/sWAWgAR1H2DcHr4PareCbAaCo1RpuU=
golang.org/x/term v0.23.0/go.mod h1:DgV24QBUrK6jhZXl+20l6UWznPlwAHm1Q1mGHtydmSk=
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
Expand Down
88 changes: 87 additions & 1 deletion pkg/client/s3client.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ package client
import (
"errors"
"fmt"
"io"
"os"
"regexp"
"sync"
Expand All @@ -31,6 +32,8 @@ import (
"github.com/IBM/ibm-cos-sdk-go/service/s3/s3manager"
"github.com/IBM/platform-services-go-sdk/resourcecontrollerv2"
"github.com/ppc64le-cloud/pvsadm/pkg"
"github.com/vbauerster/mpb"
"github.com/vbauerster/mpb/decor"
"k8s.io/klog/v2"
"k8s.io/utils/ptr"
)
Expand Down Expand Up @@ -242,11 +245,40 @@ type CustomReader struct {
mux sync.Mutex
}

type ProgressTracker struct {
progress *mpb.Progress
bar *mpb.Bar
isBarSet bool
counter *formattedCounter
}

type formattedCounter struct {
read *int64
total int64
}

// Syncable implements decor.Decorator.
func (f *formattedCounter) Syncable() (bool, chan int) {
return false, nil
}

func (f *formattedCounter) Decor(stat *decor.Statistics) string {
return fmt.Sprintf("%s / %s", formatBytes(*f.read), formatBytes(f.total))
}

func (f *formattedCounter) Format(string) (_ string, width int) {
return "", 0
}

func (f *formattedCounter) Sync() (chan int, bool) {
return nil, false
}

func (r *CustomReader) Read(p []byte) (int, error) {
return r.fp.Read(p)
}

func (r *CustomReader) ReadAt(p []byte, off int64) (int, error) {
func (r *CustomReader) ReadAt(p []byte, off int64, tracker *ProgressTracker) (int, error) {
n, err := r.fp.ReadAt(p, off)
if err != nil {
return n, err
Expand All @@ -263,6 +295,26 @@ func (r *CustomReader) ReadAt(p []byte, off int64) (int, error) {
return n, nil
}

// Format the bytes to a human-readable string
func formatBytes(size int64) string {
const (
KB = 1024
MB = KB * 1024
GB = MB * 1024
)

switch {
case size >= GB:
return fmt.Sprintf("%.2f GB", float64(size)/GB)
case size >= MB:
return fmt.Sprintf("%.2f MB", float64(size)/MB)
case size >= KB:
return fmt.Sprintf("%.2f KB", float64(size)/KB)
default:
return fmt.Sprintf("%d Bytes", size)
}
}

func (r *CustomReader) Seek(offset int64, whence int) (int64, error) {
return r.fp.Seek(offset, whence)
}
Expand All @@ -287,6 +339,40 @@ func (c *S3Client) UploadObject(fileName, objectName, bucketName string) error {
size: fileInfo.Size(),
signMap: map[int64]struct{}{},
}
progressTracker := &ProgressTracker{
progress: mpb.New(),
counter: &formattedCounter{read: new(int64), total: reader.size},
}

buffer := make([]byte, 1024)
var off int64

progressTracker.bar = progressTracker.progress.AddBar(reader.size,
mpb.PrependDecorators(
decor.Name("Uploading: ", decor.WC{W: 15}),
progressTracker.counter,
),
mpb.AppendDecorators(
decor.Percentage(),
),
)

for {
bytesRead, err := reader.ReadAt(buffer, off, progressTracker)
if err != nil {
if err == io.EOF {
break
}
klog.Errorf("An error occurred while reading file: %v", err)
break
}
if bytesRead < 1024 {
break
}
off += int64(bytesRead)
}

// klog.Infof("Operation completed successfully. Total time: %s", time.Since(start).Round(time.Second))
// Create an uploader with S3 client
uploader := s3manager.NewUploaderWithClient(c.S3Session, func(u *s3manager.Uploader) {
u.PartSize = 64 * 1024 * 1024
Expand Down

0 comments on commit 410c6d9

Please sign in to comment.