mirror of
				https://source.quilibrium.com/quilibrium/ceremonyclient.git
				synced 2025-11-04 14:47:26 +00:00 
			
		
		
		
	
		
			
				
	
	
		
			201 lines
		
	
	
		
			5.7 KiB
		
	
	
	
		
			Go
		
	
	
	
	
	
			
		
		
	
	
			201 lines
		
	
	
		
			5.7 KiB
		
	
	
	
		
			Go
		
	
	
	
	
	
package replay
 | 
						|
 | 
						|
import (
 | 
						|
	"bytes"
 | 
						|
	"crypto/rand"
 | 
						|
	"fmt"
 | 
						|
	"io"
 | 
						|
	"strconv"
 | 
						|
	"strings"
 | 
						|
	"testing"
 | 
						|
	"time"
 | 
						|
	"unicode"
 | 
						|
 | 
						|
	"github.com/cockroachdb/datadriven"
 | 
						|
	"github.com/cockroachdb/pebble"
 | 
						|
	"github.com/cockroachdb/pebble/internal/base"
 | 
						|
	"github.com/cockroachdb/pebble/internal/humanize"
 | 
						|
	"github.com/cockroachdb/pebble/vfs"
 | 
						|
	"github.com/stretchr/testify/require"
 | 
						|
)
 | 
						|
 | 
						|
func TestWorkloadCollector(t *testing.T) {
 | 
						|
	const srcDir = `src`
 | 
						|
	const destDir = `dst`
 | 
						|
	datadriven.Walk(t, "testdata/collect", func(t *testing.T, path string) {
 | 
						|
		fs := vfs.NewMem()
 | 
						|
		require.NoError(t, fs.MkdirAll(srcDir, 0755))
 | 
						|
		require.NoError(t, fs.MkdirAll(destDir, 0755))
 | 
						|
		c := NewWorkloadCollector(srcDir)
 | 
						|
		o := &pebble.Options{FS: fs}
 | 
						|
		c.Attach(o)
 | 
						|
		var currentManifest vfs.File
 | 
						|
		var buf bytes.Buffer
 | 
						|
		defer func() {
 | 
						|
			if currentManifest != nil {
 | 
						|
				currentManifest.Close()
 | 
						|
			}
 | 
						|
		}()
 | 
						|
		datadriven.RunTest(t, path, func(t *testing.T, td *datadriven.TestData) string {
 | 
						|
			buf.Reset()
 | 
						|
			switch td.Cmd {
 | 
						|
			case "cmp-files":
 | 
						|
				if len(td.CmdArgs) != 2 {
 | 
						|
					return fmt.Sprintf("expected exactly 2 args, received %d", len(td.CmdArgs))
 | 
						|
				}
 | 
						|
				b1 := readFile(t, fs, td.CmdArgs[0].String())
 | 
						|
				b2 := readFile(t, fs, td.CmdArgs[1].String())
 | 
						|
				if !bytes.Equal(b1, b2) {
 | 
						|
					return fmt.Sprintf("files are unequal: %s (%s) and %s (%s)",
 | 
						|
						td.CmdArgs[0].String(), humanize.Bytes.Uint64(uint64(len(b1))),
 | 
						|
						td.CmdArgs[1].String(), humanize.Bytes.Uint64(uint64(len(b2))))
 | 
						|
				}
 | 
						|
				return "equal"
 | 
						|
			case "clean":
 | 
						|
				for _, path := range strings.Fields(td.Input) {
 | 
						|
					typ, _, ok := base.ParseFilename(fs, path)
 | 
						|
					require.True(t, ok)
 | 
						|
					require.NoError(t, o.Cleaner.Clean(fs, typ, path))
 | 
						|
				}
 | 
						|
				return ""
 | 
						|
			case "create-manifest":
 | 
						|
				if currentManifest != nil {
 | 
						|
					require.NoError(t, currentManifest.Close())
 | 
						|
				}
 | 
						|
 | 
						|
				var fileNum uint64
 | 
						|
				var err error
 | 
						|
				td.ScanArgs(t, "filenum", &fileNum)
 | 
						|
				path := base.MakeFilepath(fs, srcDir, base.FileTypeManifest, base.DiskFileNum(fileNum))
 | 
						|
				currentManifest, err = fs.Create(path)
 | 
						|
				require.NoError(t, err)
 | 
						|
				_, err = currentManifest.Write(randData(100))
 | 
						|
				require.NoError(t, err)
 | 
						|
 | 
						|
				c.onManifestCreated(pebble.ManifestCreateInfo{
 | 
						|
					Path:    path,
 | 
						|
					FileNum: base.DiskFileNum(fileNum),
 | 
						|
				})
 | 
						|
				return ""
 | 
						|
			case "flush":
 | 
						|
				flushInfo := pebble.FlushInfo{
 | 
						|
					Done:          true,
 | 
						|
					Input:         1,
 | 
						|
					Duration:      100 * time.Millisecond,
 | 
						|
					TotalDuration: 100 * time.Millisecond,
 | 
						|
				}
 | 
						|
				for _, line := range strings.Split(td.Input, "\n") {
 | 
						|
					if line == "" {
 | 
						|
						continue
 | 
						|
					}
 | 
						|
 | 
						|
					parts := strings.FieldsFunc(line, func(r rune) bool { return unicode.IsSpace(r) || r == ':' })
 | 
						|
					tableInfo := pebble.TableInfo{Size: 10 << 10}
 | 
						|
					fileNum, err := strconv.ParseUint(parts[0], 10, 64)
 | 
						|
					require.NoError(t, err)
 | 
						|
					tableInfo.FileNum = base.FileNum(fileNum)
 | 
						|
 | 
						|
					p := writeFile(t, fs, srcDir, base.FileTypeTable, tableInfo.FileNum.DiskFileNum(), randData(int(tableInfo.Size)))
 | 
						|
					fmt.Fprintf(&buf, "created %s\n", p)
 | 
						|
					flushInfo.Output = append(flushInfo.Output, tableInfo)
 | 
						|
 | 
						|
					// Simulate a version edit applied to the current manifest.
 | 
						|
					_, err = currentManifest.Write(randData(25))
 | 
						|
					require.NoError(t, err)
 | 
						|
				}
 | 
						|
				flushInfo.InputBytes = 100 // Determinism
 | 
						|
				fmt.Fprint(&buf, flushInfo.String())
 | 
						|
				c.onFlushEnd(flushInfo)
 | 
						|
				return buf.String()
 | 
						|
			case "ingest":
 | 
						|
				ingestInfo := pebble.TableIngestInfo{}
 | 
						|
				for _, line := range strings.Split(td.Input, "\n") {
 | 
						|
					if line == "" {
 | 
						|
						continue
 | 
						|
					}
 | 
						|
 | 
						|
					parts := strings.FieldsFunc(line, func(r rune) bool { return unicode.IsSpace(r) || r == ':' })
 | 
						|
					tableInfo := pebble.TableInfo{Size: 10 << 10}
 | 
						|
					fileNum, err := strconv.ParseUint(parts[0], 10, 64)
 | 
						|
					require.NoError(t, err)
 | 
						|
					tableInfo.FileNum = base.FileNum(fileNum)
 | 
						|
 | 
						|
					p := writeFile(t, fs, srcDir, base.FileTypeTable, tableInfo.FileNum.DiskFileNum(), randData(int(tableInfo.Size)))
 | 
						|
					fmt.Fprintf(&buf, "created %s\n", p)
 | 
						|
					ingestInfo.Tables = append(ingestInfo.Tables, struct {
 | 
						|
						pebble.TableInfo
 | 
						|
						Level int
 | 
						|
					}{Level: 0, TableInfo: tableInfo})
 | 
						|
 | 
						|
					// Simulate a version edit applied to the current manifest.
 | 
						|
					_, err = currentManifest.Write(randData(25))
 | 
						|
					require.NoError(t, err)
 | 
						|
				}
 | 
						|
				fmt.Fprint(&buf, ingestInfo.String())
 | 
						|
				c.onTableIngest(ingestInfo)
 | 
						|
				return buf.String()
 | 
						|
 | 
						|
			case "ls":
 | 
						|
				return runListFiles(t, fs, td)
 | 
						|
			case "start":
 | 
						|
				c.Start(fs, destDir)
 | 
						|
				return ""
 | 
						|
			case "stat":
 | 
						|
				var buf bytes.Buffer
 | 
						|
				for _, arg := range td.CmdArgs {
 | 
						|
					fi, err := fs.Stat(arg.String())
 | 
						|
					if err != nil {
 | 
						|
						fmt.Fprintf(&buf, "%s: %s\n", arg.String(), err)
 | 
						|
						continue
 | 
						|
					}
 | 
						|
					fmt.Fprintf(&buf, "%s:\n", arg.String())
 | 
						|
					fmt.Fprintf(&buf, "  size: %d\n", fi.Size())
 | 
						|
				}
 | 
						|
				return buf.String()
 | 
						|
			case "stop":
 | 
						|
				c.Stop()
 | 
						|
				return ""
 | 
						|
			case "wait":
 | 
						|
				// Wait until all pending sstables have been copied, then list
 | 
						|
				// the files in the destination directory.
 | 
						|
				c.mu.Lock()
 | 
						|
				for c.mu.tablesEnqueued != c.mu.tablesCopied {
 | 
						|
					c.mu.copyCond.Wait()
 | 
						|
				}
 | 
						|
				c.mu.Unlock()
 | 
						|
				listFiles(t, fs, &buf, destDir)
 | 
						|
				return buf.String()
 | 
						|
			default:
 | 
						|
				return fmt.Sprintf("unrecognized command %q", td.Cmd)
 | 
						|
			}
 | 
						|
		})
 | 
						|
	})
 | 
						|
}
 | 
						|
 | 
						|
func randData(byteCount int) []byte {
 | 
						|
	b := make([]byte, byteCount)
 | 
						|
	rand.Read(b)
 | 
						|
	return b
 | 
						|
}
 | 
						|
 | 
						|
func writeFile(
 | 
						|
	t *testing.T, fs vfs.FS, dir string, typ base.FileType, fileNum base.DiskFileNum, data []byte,
 | 
						|
) string {
 | 
						|
	path := base.MakeFilepath(fs, dir, typ, fileNum)
 | 
						|
	f, err := fs.Create(path)
 | 
						|
	require.NoError(t, err)
 | 
						|
	_, err = f.Write(data)
 | 
						|
	require.NoError(t, err)
 | 
						|
	require.NoError(t, f.Close())
 | 
						|
	return path
 | 
						|
}
 | 
						|
 | 
						|
func readFile(t *testing.T, fs vfs.FS, path string) []byte {
 | 
						|
	r, err := fs.Open(path)
 | 
						|
	require.NoError(t, err)
 | 
						|
	b, err := io.ReadAll(r)
 | 
						|
	require.NoError(t, err)
 | 
						|
	require.NoError(t, r.Close())
 | 
						|
	return b
 | 
						|
}
 |