package storage import ( "bytes" "crypto/rand" "io" "os" "testing" "github.com/docker/distribution/context" "github.com/docker/distribution/digest" storagedriver "github.com/docker/distribution/registry/storage/driver" "github.com/docker/distribution/registry/storage/driver/inmemory" ) // TestSimpleWrite takes the fileWriter through common write operations // ensuring data integrity. func TestSimpleWrite(t *testing.T) { content := make([]byte, 1<<20) n, err := rand.Read(content) if err != nil { t.Fatalf("unexpected error building random data: %v", err) } if n != len(content) { t.Fatalf("random read did't fill buffer") } dgst, err := digest.FromReader(bytes.NewReader(content)) if err != nil { t.Fatalf("unexpected error digesting random content: %v", err) } driver := inmemory.New() path := "/random" ctx := context.Background() fw, err := newFileWriter(ctx, driver, path) if err != nil { t.Fatalf("unexpected error creating fileWriter: %v", err) } defer fw.Close() n, err = fw.Write(content) if err != nil { t.Fatalf("unexpected error writing content: %v", err) } fw.Flush() if n != len(content) { t.Fatalf("unexpected write length: %d != %d", n, len(content)) } fr, err := newFileReader(ctx, driver, path) if err != nil { t.Fatalf("unexpected error creating fileReader: %v", err) } defer fr.Close() verifier, err := digest.NewDigestVerifier(dgst) if err != nil { t.Fatalf("unexpected error getting digest verifier: %s", err) } io.Copy(verifier, fr) if !verifier.Verified() { t.Fatalf("unable to verify write data") } // Check the seek position is equal to the content length end, err := fw.Seek(0, os.SEEK_END) if err != nil { t.Fatalf("unexpected error seeking: %v", err) } if end != int64(len(content)) { t.Fatalf("write did not advance offset: %d != %d", end, len(content)) } // Double the content, but use the WriteAt method doubled := append(content, content...) doubledgst, err := digest.FromReader(bytes.NewReader(doubled)) if err != nil { t.Fatalf("unexpected error digesting doubled content: %v", err) } n, err = fw.WriteAt(content, end) if err != nil { t.Fatalf("unexpected error writing content at %d: %v", end, err) } if n != len(content) { t.Fatalf("writeat was short: %d != %d", n, len(content)) } fr, err = newFileReader(ctx, driver, path) if err != nil { t.Fatalf("unexpected error creating fileReader: %v", err) } defer fr.Close() verifier, err = digest.NewDigestVerifier(doubledgst) if err != nil { t.Fatalf("unexpected error getting digest verifier: %s", err) } io.Copy(verifier, fr) if !verifier.Verified() { t.Fatalf("unable to verify write data") } // Check that WriteAt didn't update the offset. end, err = fw.Seek(0, os.SEEK_END) if err != nil { t.Fatalf("unexpected error seeking: %v", err) } if end != int64(len(content)) { t.Fatalf("write did not advance offset: %d != %d", end, len(content)) } // Now, we copy from one path to another, running the data through the // fileReader to fileWriter, rather than the driver.Move command to ensure // everything is working correctly. fr, err = newFileReader(ctx, driver, path) if err != nil { t.Fatalf("unexpected error creating fileReader: %v", err) } defer fr.Close() fw, err = newFileWriter(ctx, driver, "/copied") if err != nil { t.Fatalf("unexpected error creating fileWriter: %v", err) } defer fw.Close() nn, err := io.Copy(fw, fr) if err != nil { t.Fatalf("unexpected error copying data: %v", err) } if nn != int64(len(doubled)) { t.Fatalf("unexpected copy length: %d != %d", nn, len(doubled)) } fr, err = newFileReader(ctx, driver, "/copied") if err != nil { t.Fatalf("unexpected error creating fileReader: %v", err) } defer fr.Close() verifier, err = digest.NewDigestVerifier(doubledgst) if err != nil { t.Fatalf("unexpected error getting digest verifier: %s", err) } io.Copy(verifier, fr) if !verifier.Verified() { t.Fatalf("unable to verify write data") } } func TestBufferedFileWriter(t *testing.T) { ctx := context.Background() writer, err := newFileWriter(ctx, inmemory.New(), "/random") if err != nil { t.Fatalf("Failed to initialize bufferedFileWriter: %v", err.Error()) } // write one byte and ensure the offset hasn't been incremented. // offset will only get incremented when the buffer gets flushed short := []byte{byte(1)} writer.Write(short) if writer.offset > 0 { t.Fatalf("WriteStream called prematurely") } // write enough data to cause the buffer to flush and confirm // the offset has been incremented long := make([]byte, fileWriterBufferSize) _, err = rand.Read(long) if err != nil { t.Fatalf("unexpected error building random data: %v", err) } for i := range long { long[i] = byte(i) } writer.Write(long) writer.Close() if writer.offset != (fileWriterBufferSize + 1) { t.Fatalf("WriteStream not called when buffer capacity reached") } } func BenchmarkFileWriter(b *testing.B) { b.StopTimer() // not sure how long setup above will take for i := 0; i < b.N; i++ { // Start basic fileWriter initialization fw := fileWriter{ driver: inmemory.New(), path: "/random", } ctx := context.Background() if fi, err := fw.driver.Stat(ctx, fw.path); err != nil { switch err := err.(type) { case storagedriver.PathNotFoundError: // ignore, offset is zero default: b.Fatalf("Failed to initialize fileWriter: %v", err.Error()) } } else { if fi.IsDir() { b.Fatalf("Cannot write to a directory") } fw.size = fi.Size() } randomBytes := make([]byte, 1<<20) _, err := rand.Read(randomBytes) if err != nil { b.Fatalf("unexpected error building random data: %v", err) } // End basic file writer initialization b.StartTimer() for j := 0; j < 100; j++ { fw.Write(randomBytes) } b.StopTimer() } } func BenchmarkBufferedFileWriter(b *testing.B) { b.StopTimer() // not sure how long setup above will take ctx := context.Background() for i := 0; i < b.N; i++ { bfw, err := newFileWriter(ctx, inmemory.New(), "/random") if err != nil { b.Fatalf("Failed to initialize bufferedFileWriter: %v", err.Error()) } randomBytes := make([]byte, 1<<20) _, err = rand.Read(randomBytes) if err != nil { b.Fatalf("unexpected error building random data: %v", err) } b.StartTimer() for j := 0; j < 100; j++ { bfw.Write(randomBytes) } b.StopTimer() } }