kodotesting.go

kodotesting.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"os/user"
	"path"
	"runtime"
	"time"

	"github.com/dolab/gogo"
	"github.com/golib/cli"
	"qiniupkg.com/api.v7/kodo"
)

// TODO: add evmclient for authed request!
var (
	kodologger gogo.Logger
	kodoconfig kodo.Config
	kodoclient *kodo.Client

	tryHomeDir = func() string {
		if runtime.GOOS == "windows" {
			home := os.Getenv("HOMEDRIVE") + os.Getenv("HOMEPATH")
			if home == "" {
				home = os.Getenv("USERPROFILE")
			}

			return home
		}

		return os.Getenv("HOME")
	}
)

func main() {
	kodologger = gogo.NewAppLogger("stderr", "")
	kodologger.SetColor(true)

	cu, err := user.Current()
	if err != nil {
		kodologger.Errorf("user.Current(): %v", err)

		cu = &user.User{
			HomeDir: tryHomeDir(),
		}
	}

	// resolve kodo config
	filename := path.Join(cu.HomeDir, ".kodotesting", "config.json")
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		kodologger.Errorf("ioutil.ReadFile(%s): %v", filename, err)
		return
	}

	err = json.Unmarshal(data, &kodoconfig)
	if err != nil {
		kodologger.Errorf("json.Unmarshal(%T): %v", kodoconfig, err)
		return
	}

	// setup kodo sdk
	kodoclient = kodo.NewWithoutZone(&kodoconfig)

	// new app
	app := cli.NewApp()
	app.Name = "kodotesting"
	app.Usage = "a new qiniu evm to kodo speed testing commandline utils"
	app.Version = "1.0.0"
	app.Authors = []cli.Author{
		{
			Name:  "Spring MC",
			Email: "[email protected]",
		},
	}

	app.Commands = []cli.Command{
		{
			Name:  "upload",
			Usage: "upload file to kodo with speed testing, see kodotesting upload",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "bucket",
					Usage: "specify bucket name for testing",
				},
				cli.StringFlag{
					Name:  "file",
					Usage: "specify abs path of testing file",
				},
			},
			Action: func(ctx *cli.Context) error {
				bucket := ctx.String("bucket")
				file := ctx.String("file")
				if bucket == "" || file == "" {
					kodologger.Warn("Both bucket and file are required!")

					cli.ShowCommandHelp(ctx, "upload")
					return nil
				}

				file = path.Clean(file)

				fstat, ferr := os.Stat(file)
				if ferr != nil {
					kodologger.Errorf("os.Stat(%s): %v", file, ferr)

					return ferr
				}

				startedAt := time.Now()
				kodologger.Infof("Starting kodo.PutFile(%s, %s) at %v", bucket, file, startedAt)

				err := kodoclient.Bucket(bucket).PutFile(context.Background(), nil, path.Base(file), file, nil)
				if err != nil {
					kodologger.Errorf("kodo.PutFile(%s, %s): %v", bucket, file, err)
				} else {
					since := time.Since(startedAt)
					kodologger.Infof("Finished kodo.PutFile(%s, %s, %d) in %v (speed %v)", bucket, file, fstat.Size(), since, speedHumanize(float64(fstat.Size())/(float64(since)/float64(time.Second))))
				}

				return err
			},
		},
	}

	app.Run(os.Args)
}

func speedHumanize(speed float64) string {
	var (
		gb = float64(1 << 30)
		mb = float64(1 << 20)
		kb = float64(1 << 10)
	)

	if speed > gb {
		return fmt.Sprintf("%.2f G/s", (speed / gb))
	}

	if speed > mb {
		return fmt.Sprintf("%.2f M/s", (speed / mb))
	}

	if speed > kb {
		return fmt.Sprintf("%.2f K/s", (speed / kb))
	}

	return fmt.Sprintf("%.2f B/s", speed)
}