github.com/fawick/restic@v0.1.1-0.20171126184616-c02923fbfc79/internal/repository/packer_manager_test.go (about)

     1  package repository
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"math/rand"
     7  	"os"
     8  	"testing"
     9  
    10  	"github.com/restic/restic/internal/backend/mem"
    11  	"github.com/restic/restic/internal/crypto"
    12  	"github.com/restic/restic/internal/fs"
    13  	"github.com/restic/restic/internal/mock"
    14  	"github.com/restic/restic/internal/restic"
    15  )
    16  
    17  type randReader struct {
    18  	src  rand.Source
    19  	rand *rand.Rand
    20  }
    21  
    22  func newRandReader(src rand.Source) *randReader {
    23  	return &randReader{
    24  		src:  src,
    25  		rand: rand.New(src),
    26  	}
    27  }
    28  
    29  // Read generates len(p) random bytes and writes them into p. It
    30  // always returns len(p) and a nil error.
    31  func (r *randReader) Read(p []byte) (n int, err error) {
    32  	for i := 0; i < len(p); i += 7 {
    33  		val := r.src.Int63()
    34  		for j := 0; i+j < len(p) && j < 7; j++ {
    35  			p[i+j] = byte(val)
    36  			val >>= 8
    37  		}
    38  	}
    39  	return len(p), nil
    40  }
    41  
    42  func randomID(rd io.Reader) restic.ID {
    43  	id := restic.ID{}
    44  	_, err := io.ReadFull(rd, id[:])
    45  	if err != nil {
    46  		panic(err)
    47  	}
    48  	return id
    49  }
    50  
    51  const maxBlobSize = 1 << 20
    52  
    53  func saveFile(t testing.TB, be Saver, f *os.File, id restic.ID) {
    54  	h := restic.Handle{Type: restic.DataFile, Name: id.String()}
    55  	t.Logf("save file %v", h)
    56  
    57  	if err := be.Save(context.TODO(), h, f); err != nil {
    58  		t.Fatal(err)
    59  	}
    60  
    61  	if err := f.Close(); err != nil {
    62  		t.Fatal(err)
    63  	}
    64  
    65  	if err := fs.RemoveIfExists(f.Name()); err != nil {
    66  		t.Fatal(err)
    67  	}
    68  }
    69  
    70  func fillPacks(t testing.TB, rnd *randReader, be Saver, pm *packerManager, buf []byte) (bytes int) {
    71  	for i := 0; i < 100; i++ {
    72  		l := rnd.rand.Intn(1 << 20)
    73  		seed := rnd.rand.Int63()
    74  
    75  		packer, err := pm.findPacker()
    76  		if err != nil {
    77  			t.Fatal(err)
    78  		}
    79  
    80  		rd := newRandReader(rand.NewSource(seed))
    81  		id := randomID(rd)
    82  		buf = buf[:l]
    83  		_, err = io.ReadFull(rd, buf)
    84  		if err != nil {
    85  			t.Fatal(err)
    86  		}
    87  
    88  		n, err := packer.Add(restic.DataBlob, id, buf)
    89  		if n != l {
    90  			t.Errorf("Add() returned invalid number of bytes: want %v, got %v", n, l)
    91  		}
    92  		bytes += l
    93  
    94  		if packer.Size() < minPackSize {
    95  			pm.insertPacker(packer)
    96  			continue
    97  		}
    98  
    99  		_, err = packer.Finalize()
   100  		if err != nil {
   101  			t.Fatal(err)
   102  		}
   103  
   104  		if _, err = packer.tmpfile.Seek(0, 0); err != nil {
   105  			t.Fatal(err)
   106  		}
   107  
   108  		packID := restic.IDFromHash(packer.hw.Sum(nil))
   109  		saveFile(t, be, packer.tmpfile, packID)
   110  	}
   111  
   112  	return bytes
   113  }
   114  
   115  func flushRemainingPacks(t testing.TB, rnd *randReader, be Saver, pm *packerManager) (bytes int) {
   116  	if pm.countPacker() > 0 {
   117  		for _, packer := range pm.packers {
   118  			n, err := packer.Finalize()
   119  			if err != nil {
   120  				t.Fatal(err)
   121  			}
   122  			bytes += int(n)
   123  
   124  			packID := restic.IDFromHash(packer.hw.Sum(nil))
   125  			saveFile(t, be, packer.tmpfile, packID)
   126  		}
   127  	}
   128  
   129  	return bytes
   130  }
   131  
   132  func TestPackerManager(t *testing.T) {
   133  	rnd := newRandReader(rand.NewSource(23))
   134  
   135  	be := mem.New()
   136  	pm := newPackerManager(be, crypto.NewRandomKey())
   137  
   138  	blobBuf := make([]byte, maxBlobSize)
   139  
   140  	bytes := fillPacks(t, rnd, be, pm, blobBuf)
   141  	bytes += flushRemainingPacks(t, rnd, be, pm)
   142  
   143  	t.Logf("saved %d bytes", bytes)
   144  }
   145  
   146  func BenchmarkPackerManager(t *testing.B) {
   147  	rnd := newRandReader(rand.NewSource(23))
   148  
   149  	be := &mock.Backend{
   150  		SaveFn: func(context.Context, restic.Handle, io.Reader) error { return nil },
   151  	}
   152  	blobBuf := make([]byte, maxBlobSize)
   153  
   154  	t.ResetTimer()
   155  
   156  	for i := 0; i < t.N; i++ {
   157  		bytes := 0
   158  		pm := newPackerManager(be, crypto.NewRandomKey())
   159  		bytes += fillPacks(t, rnd, be, pm, blobBuf)
   160  		bytes += flushRemainingPacks(t, rnd, be, pm)
   161  		t.Logf("saved %d bytes", bytes)
   162  	}
   163  }