lockedfile.go

     1  // Copyright 2018 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // Package lockedfile creates and manipulates files whose contents should only
     6  // change atomically.
     7  package lockedfile
     8  
     9  import (
    10  	"fmt"
    11  	"io"
    12  	"io/fs"
    13  	"os"
    14  	"runtime"
    15  )
    16  
    17  // A File is a locked *os.File.
    18  //
    19  // Closing the file releases the lock.
    20  //
    21  // If the program exits while a file is locked, the operating system releases
    22  // the lock but may not do so promptly: callers must ensure that all locked
    23  // files are closed before exiting.
    24  type File struct {
    25  	osFile
    26  	closed bool
    27  	// cleanup panics when the file is no longer referenced and it has not been closed.
    28  	cleanup runtime.Cleanup
    29  }
    30  
    31  // osFile embeds a *os.File while keeping the pointer itself unexported.
    32  // (When we close a File, it must be the same file descriptor that we opened!)
    33  type osFile struct {
    34  	*os.File
    35  }
    36  
    37  // OpenFile is like os.OpenFile, but returns a locked file.
    38  // If flag includes os.O_WRONLY or os.O_RDWR, the file is write-locked;
    39  // otherwise, it is read-locked.
    40  func OpenFile(name string, flag int, perm fs.FileMode) (*File, error) {
    41  	var (
    42  		f   = new(File)
    43  		err error
    44  	)
    45  	f.osFile.File, err = openFile(name, flag, perm)
    46  	if err != nil {
    47  		return nil, err
    48  	}
    49  
    50  	// Although the operating system will drop locks for open files when the go
    51  	// command exits, we want to hold locks for as little time as possible, and we
    52  	// especially don't want to leave a file locked after we're done with it. Our
    53  	// Close method is what releases the locks, so use a cleanup to report
    54  	// missing Close calls on a best-effort basis.
    55  	f.cleanup = runtime.AddCleanup(f, func(fileName string) {
    56  		panic(fmt.Sprintf("lockedfile.File %s became unreachable without a call to Close", fileName))
    57  	}, f.Name())
    58  
    59  	return f, nil
    60  }
    61  
    62  // Open is like os.Open, but returns a read-locked file.
    63  func Open(name string) (*File, error) {
    64  	return OpenFile(name, os.O_RDONLY, 0)
    65  }
    66  
    67  // Create is like os.Create, but returns a write-locked file.
    68  func Create(name string) (*File, error) {
    69  	return OpenFile(name, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0666)
    70  }
    71  
    72  // Edit creates the named file with mode 0666 (before umask),
    73  // but does not truncate existing contents.
    74  //
    75  // If Edit succeeds, methods on the returned File can be used for I/O.
    76  // The associated file descriptor has mode O_RDWR and the file is write-locked.
    77  func Edit(name string) (*File, error) {
    78  	return OpenFile(name, os.O_RDWR|os.O_CREATE, 0666)
    79  }
    80  
    81  // Close unlocks and closes the underlying file.
    82  //
    83  // Close may be called multiple times; all calls after the first will return a
    84  // non-nil error.
    85  func (f *File) Close() error {
    86  	if f.closed {
    87  		return &fs.PathError{
    88  			Op:   "close",
    89  			Path: f.Name(),
    90  			Err:  fs.ErrClosed,
    91  		}
    92  	}
    93  	f.closed = true
    94  
    95  	err := closeFile(f.osFile.File)
    96  	f.cleanup.Stop()
    97  	// f may be dead at the moment after we access f.cleanup,
    98  	// so the cleanup can fire before Stop completes. Keep f
    99  	// alive while we call Stop. See the documentation for
   100  	// runtime.Cleanup.Stop.
   101  	runtime.KeepAlive(f)
   102  	return err
   103  }
   104  
   105  // Read opens the named file with a read-lock and returns its contents.
   106  func Read(name string) ([]byte, error) {
   107  	f, err := Open(name)
   108  	if err != nil {
   109  		return nil, err
   110  	}
   111  	defer f.Close()
   112  
   113  	return io.ReadAll(f)
   114  }
   115  
   116  // Write opens the named file (creating it with the given permissions if needed),
   117  // then write-locks it and overwrites it with the given content.
   118  func Write(name string, content io.Reader, perm fs.FileMode) (err error) {
   119  	f, err := OpenFile(name, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, perm)
   120  	if err != nil {
   121  		return err
   122  	}
   123  
   124  	_, err = io.Copy(f, content)
   125  	if closeErr := f.Close(); err == nil {
   126  		err = closeErr
   127  	}
   128  	return err
   129  }
   130  
   131  // Transform invokes t with the result of reading the named file, with its lock
   132  // still held.
   133  //
   134  // If t returns a nil error, Transform then writes the returned contents back to
   135  // the file, making a best effort to preserve existing contents on error.
   136  //
   137  // t must not modify the slice passed to it.
   138  func Transform(name string, t func([]byte) ([]byte, error)) (err error) {
   139  	f, err := Edit(name)
   140  	if err != nil {
   141  		return err
   142  	}
   143  	defer f.Close()
   144  
   145  	old, err := io.ReadAll(f)
   146  	if err != nil {
   147  		return err
   148  	}
   149  
   150  	new, err := t(old)
   151  	if err != nil {
   152  		return err
   153  	}
   154  
   155  	if len(new) > len(old) {
   156  		// The overall file size is increasing, so write the tail first: if we're
   157  		// about to run out of space on the disk, we would rather detect that
   158  		// failure before we have overwritten the original contents.
   159  		if _, err := f.WriteAt(new[len(old):], int64(len(old))); err != nil {
   160  			// Make a best effort to remove the incomplete tail.
   161  			f.Truncate(int64(len(old)))
   162  			return err
   163  		}
   164  	}
   165  
   166  	// We're about to overwrite the old contents. In case of failure, make a best
   167  	// effort to roll back before we close the file.
   168  	defer func() {
   169  		if err != nil {
   170  			if _, err := f.WriteAt(old, 0); err == nil {
   171  				f.Truncate(int64(len(old)))
   172  			}
   173  		}
   174  	}()
   175  
   176  	if len(new) >= len(old) {
   177  		if _, err := f.WriteAt(new[:len(old)], 0); err != nil {
   178  			return err
   179  		}
   180  	} else {
   181  		if _, err := f.WriteAt(new, 0); err != nil {
   182  			return err
   183  		}
   184  		// The overall file size is decreasing, so shrink the file to its final size
   185  		// after writing. We do this after writing (instead of before) so that if
   186  		// the write fails, enough filesystem space will likely still be reserved
   187  		// to contain the previous contents.
   188  		if err := f.Truncate(int64(len(new))); err != nil {
   189  			return err
   190  		}
   191  	}
   192  
   193  	return nil
   194  }
   195
View as plain text