test/iotest/iotest.go

package main

import (
	"bytes"
	"encoding/base64"
	"flag"
	"fmt"
	"io"
	"math"
	"math/rand"
	"os"
	"path"
	"path/filepath"
	"sync/atomic"
	"time"

	"git.wecise.com/wecise/common/cassandra"
	"git.wecise.com/wecise/common/dfs"
	"git.wecise.com/wecise/common/etcd"
	"git.wecise.com/wecise/common/matrix/logger"
	"gopkg.in/src-d/go-billy.v4"
)

// 一个持续并发文件读写验证的程序
// 在指定目录下写入一些随机内容的文本文件
// 写入完成后读取并验证写入文件内容的正确性
// 验证正确后在指定的延时时间后将文件删除
// 文件尺寸可以在指定的范围内随机生成
// 所有可指定变量以命令行参数形式传入，并设有默认值
// 每秒在控制台输出验证成功的文件计数信息
// 验证内容不一致时保留文件不删除，提示验证失败的文件路径，程序结束。

func init() {
	logger.SetFormat("yyyy-MM-dd HH:mm:ss [pid] [level] msg", "\n")
	logger.SetRollingFile("", "iotest.log", -1, 1*1024*1024, math.MaxInt64, 1)
}

var mdfs billy.Filesystem

func initDFS(keyspace string) error {
	cassandraOption := cassandra.Option{Keyspace: keyspace, DisableInitialHostLookup: os.Getenv("VERSION") == "dev"}
	ccli, err := cassandra.GetClientByOption(cassandraOption)
	if err != nil {
		return fmt.Errorf("cassandra.GetClient Error: %v", err)
	}
	ecli, err := etcd.Get()
	if err != nil {
		return fmt.Errorf("etcd.Get Error: %v", err)
	}
	mdfs = dfs.New(keyspace, ccli.Session, ecli)
	return nil
}

// func writefile(filePath string, content []byte) (err error) {
// 	dir := path.Dir(filePath)
// 	f, e := os.Stat(dir)
// 	if os.IsNotExist(e) || !f.IsDir() {
// 		if err = os.MkdirAll(dir, os.ModePerm); err != nil {
// 			return
// 		}
// 	}
// 	// 写入文件
// 	err = ioutil.WriteFile(filePath, content, 0777)
// 	return
// }

// func readfile(filePath string) ([]byte, error) {
// 	return ioutil.ReadFile(filePath)
// }

// func deletefile(filePath string) error {
// 	return os.Remove(filePath)
// }

func writefile(filePath string, content []byte) error {
	dir := path.Dir(filePath)
	fi, e := mdfs.Stat(dir)
	if os.IsNotExist(e) || !fi.IsDir() {
		if err := mdfs.MkdirAll(dir, os.ModePerm); err != nil {
			return err
		}
	}
	// 写入文件
	f, err := mdfs.OpenFile(filePath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, os.ModePerm)
	if err != nil {
		return err
	}
	_, err = f.Write(content)
	if err1 := f.Close(); err1 != nil && err == nil {
		err = err1
	}
	return err
}

func readfile(filePath string) ([]byte, error) {
	var size int
	if info, err := mdfs.Stat(filePath); err == nil {
		size64 := info.Size()
		if int64(int(size64)) == size64 {
			size = int(size64)
		}
	}

	f, err := mdfs.Open(filePath)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	size++ // one byte for final read at EOF

	// If a file claims a small size, read at least 512 bytes.
	// In particular, files in Linux's /proc claim size 0 but
	// then do not work right if read in small pieces,
	// so an initial read of 1 byte would not work correctly.
	if size < 512 {
		size = 512
	}

	data := make([]byte, 0, size)
	for {
		if len(data) >= cap(data) {
			d := append(data[:cap(data)], 0)
			data = d[:len(data)]
		}
		n, err := f.Read(data[len(data):cap(data)])
		data = data[:len(data)+n]
		if err != nil {
			if err == io.EOF {
				err = nil
			}
			return data, err
		}
	}
}

func deletefile(filePath string) error {
	return mdfs.Remove(filePath)
}

func main() {
	// 设置命令行参数
	keyspace := flag.String("keyspace", "ootest", "keyspace")
	rootPath := flag.String("root-path", "_test", "文件目录")
	delay := flag.Duration("delay", 30*time.Second, "删除延迟时间")
	concurLimit := flag.Int("concur-limit", 10, "最大并发数")
	fileSizeMin := flag.Int("file-size-min", 1024, "最小文件尺寸")
	fileSizeMax := flag.Int("file-size-max", 10240, "最大文件尺寸")
	flag.Parse()

	// 初始化随机数生成器
	rand.Seed(time.Now().UnixNano())

	if err := initDFS(*keyspace); err != nil {
		logger.Error(err)
	}

	type fileinfo struct {
		filepath string
		filesize int
		checksum uint64
		content  string
	}
	checkChan := make(chan *fileinfo, *concurLimit)

	logger.Info("开始验证")

	// 并发写入文件
	go func() {
		writeChan := make(chan struct{}, *concurLimit)
		for i := 0; ; i++ {
			// 并发控制
			writeChan <- struct{}{}
			go func(i int) {
				defer func() {
					<-writeChan
				}()
				// 随机生成文件名和文件内容
				dirName := fmt.Sprintf("%02d", rand.Intn(99)+1)
				fileName := fmt.Sprintf("%06d", i) + ".txt"
				filePath := filepath.Join(*rootPath, dirName, fileName)
				fileBinarySize := (rand.Intn(*fileSizeMax-*fileSizeMin) + *fileSizeMin) / 4 * 3 // 二进制长度

				// 随机生成文件名和文件内容
				bs := make([]byte, fileBinarySize)
				rand.Read(bs)
				content := base64.RawURLEncoding.EncodeToString(bs)

				// 写入文件
				if err := writefile(filePath, []byte(content)); err != nil {
					logger.Error("写入文件", filePath, "失败：", err)
					return
				}

				// 写入完成
				checkChan <- &fileinfo{
					filepath: filePath,
					filesize: len(content),
					content:  content,
				}
			}(i)
		}
	}()

	var successCount int32

	// 并发验证和删除文件
	go func() {
		for fi := range checkChan {
			go func(fi *fileinfo) {
				// 读取文件
				content, err := readfile(fi.filepath)
				if err != nil {
					logger.Error("读取文件", fi.filepath, "失败：", err)
					return
				}

				// 验证文件内容是否正确
				expectedSize := fi.filesize
				if len(content) != expectedSize || !bytes.Equal(content, []byte(fi.content)) {
					logger.Error("验证失败", fi.filepath, fmt.Sprintf("文件内容与期望值不一致，写入内容为：\n%s\n读出内容为：\n%s", string(content), fi.content))
					os.Exit(1)
					return
				}

				// 验证成功，统计计数
				atomic.AddInt32(&successCount, 1)

				// 延迟删除文件
				time.AfterFunc(*delay, func() {
					if err := deletefile(fi.filepath); err != nil {
						logger.Error("删除文件", fi.filepath, "失败：", err)
					}
				})
			}(fi)
		}
	}()

	t := time.NewTicker(2 * time.Second)
	for {
		select {
		case <-t.C:
			// 输出计数信息
			logger.Info("验证成功的文件数：", successCount)
		}
	}
}