github.com/mika/distribution@v2.2.2-0.20160108133430-a75790e3d8e0+incompatible/registry/storage/filewriter.go (about)

     1  package storage
     2  
     3  import (
     4  	"bufio"
     5  	"bytes"
     6  	"fmt"
     7  	"io"
     8  	"os"
     9  
    10  	"github.com/docker/distribution/context"
    11  	storagedriver "github.com/docker/distribution/registry/storage/driver"
    12  )
    13  
    14  const (
    15  	fileWriterBufferSize = 5 << 20
    16  )
    17  
    18  // fileWriter implements a remote file writer backed by a storage driver.
    19  type fileWriter struct {
    20  	driver storagedriver.StorageDriver
    21  
    22  	ctx context.Context
    23  
    24  	// identifying fields
    25  	path string
    26  
    27  	// mutable fields
    28  	size   int64 // size of the file, aka the current end
    29  	offset int64 // offset is the current write offset
    30  	err    error // terminal error, if set, reader is closed
    31  }
    32  
    33  type bufferedFileWriter struct {
    34  	fileWriter
    35  	bw *bufio.Writer
    36  }
    37  
    38  // fileWriterInterface makes the desired io compliant interface that the
    39  // filewriter should implement.
    40  type fileWriterInterface interface {
    41  	io.WriteSeeker
    42  	io.ReaderFrom
    43  	io.Closer
    44  }
    45  
    46  var _ fileWriterInterface = &fileWriter{}
    47  
    48  // newFileWriter returns a prepared fileWriter for the driver and path. This
    49  // could be considered similar to an "open" call on a regular filesystem.
    50  func newFileWriter(ctx context.Context, driver storagedriver.StorageDriver, path string) (*bufferedFileWriter, error) {
    51  	fw := fileWriter{
    52  		driver: driver,
    53  		path:   path,
    54  		ctx:    ctx,
    55  	}
    56  
    57  	if fi, err := driver.Stat(ctx, path); err != nil {
    58  		switch err := err.(type) {
    59  		case storagedriver.PathNotFoundError:
    60  			// ignore, offset is zero
    61  		default:
    62  			return nil, err
    63  		}
    64  	} else {
    65  		if fi.IsDir() {
    66  			return nil, fmt.Errorf("cannot write to a directory")
    67  		}
    68  
    69  		fw.size = fi.Size()
    70  	}
    71  
    72  	buffered := bufferedFileWriter{
    73  		fileWriter: fw,
    74  	}
    75  	buffered.bw = bufio.NewWriterSize(&buffered.fileWriter, fileWriterBufferSize)
    76  
    77  	return &buffered, nil
    78  }
    79  
    80  // wraps the fileWriter.Write method to buffer small writes
    81  func (bfw *bufferedFileWriter) Write(p []byte) (int, error) {
    82  	return bfw.bw.Write(p)
    83  }
    84  
    85  // wraps fileWriter.Close to ensure the buffer is flushed
    86  // before we close the writer.
    87  func (bfw *bufferedFileWriter) Close() (err error) {
    88  	if err = bfw.Flush(); err != nil {
    89  		return err
    90  	}
    91  	err = bfw.fileWriter.Close()
    92  	return err
    93  }
    94  
    95  // wraps fileWriter.Seek to ensure offset is handled
    96  // correctly in respect to pending data in the buffer
    97  func (bfw *bufferedFileWriter) Seek(offset int64, whence int) (int64, error) {
    98  	if err := bfw.Flush(); err != nil {
    99  		return 0, err
   100  	}
   101  	return bfw.fileWriter.Seek(offset, whence)
   102  }
   103  
   104  // wraps bufio.Writer.Flush to allow intermediate flushes
   105  // of the bufferedFileWriter
   106  func (bfw *bufferedFileWriter) Flush() error {
   107  	return bfw.bw.Flush()
   108  }
   109  
   110  // Write writes the buffer p at the current write offset.
   111  func (fw *fileWriter) Write(p []byte) (n int, err error) {
   112  	nn, err := fw.ReadFrom(bytes.NewReader(p))
   113  	return int(nn), err
   114  }
   115  
   116  // ReadFrom reads reader r until io.EOF writing the contents at the current
   117  // offset.
   118  func (fw *fileWriter) ReadFrom(r io.Reader) (n int64, err error) {
   119  	if fw.err != nil {
   120  		return 0, fw.err
   121  	}
   122  
   123  	nn, err := fw.driver.WriteStream(fw.ctx, fw.path, fw.offset, r)
   124  
   125  	// We should forward the offset, whether or not there was an error.
   126  	// Basically, we keep the filewriter in sync with the reader's head. If an
   127  	// error is encountered, the whole thing should be retried but we proceed
   128  	// from an expected offset, even if the data didn't make it to the
   129  	// backend.
   130  	fw.offset += nn
   131  
   132  	if fw.offset > fw.size {
   133  		fw.size = fw.offset
   134  	}
   135  
   136  	return nn, err
   137  }
   138  
   139  // Seek moves the write position do the requested offest based on the whence
   140  // argument, which can be os.SEEK_CUR, os.SEEK_END, or os.SEEK_SET.
   141  func (fw *fileWriter) Seek(offset int64, whence int) (int64, error) {
   142  	if fw.err != nil {
   143  		return 0, fw.err
   144  	}
   145  
   146  	var err error
   147  	newOffset := fw.offset
   148  
   149  	switch whence {
   150  	case os.SEEK_CUR:
   151  		newOffset += int64(offset)
   152  	case os.SEEK_END:
   153  		newOffset = fw.size + int64(offset)
   154  	case os.SEEK_SET:
   155  		newOffset = int64(offset)
   156  	}
   157  
   158  	if newOffset < 0 {
   159  		err = fmt.Errorf("cannot seek to negative position")
   160  	} else {
   161  		// No problems, set the offset.
   162  		fw.offset = newOffset
   163  	}
   164  
   165  	return fw.offset, err
   166  }
   167  
   168  // Close closes the fileWriter for writing.
   169  // Calling it once is valid and correct and it will
   170  // return a nil error. Calling it subsequent times will
   171  // detect that fw.err has been set and will return the error.
   172  func (fw *fileWriter) Close() error {
   173  	if fw.err != nil {
   174  		return fw.err
   175  	}
   176  
   177  	fw.err = fmt.Errorf("filewriter@%v: closed", fw.path)
   178  
   179  	return nil
   180  }