github.com/grailbio/base@v0.0.11/state/file.go (about)

     1  // Copyright 2018 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache-2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  // Package state implements atomic file-based state management
     6  // with support for advisory locking.
     7  package state
     8  
     9  import (
    10  	"encoding/json"
    11  	"errors"
    12  	"io/ioutil"
    13  	"os"
    14  	"path/filepath"
    15  	"sync"
    16  	"syscall"
    17  )
    18  
    19  // ErrNoState is returned when attempting to read a nonexistent state.
    20  var ErrNoState = errors.New("no state exists")
    21  
    22  // File implements file-based state management with support for
    23  // advisory locking. It is also safe to use concurrently within a
    24  // process.
    25  type File struct {
    26  	mu     sync.Mutex
    27  	prefix string
    28  	lockfd int
    29  }
    30  
    31  // New creates and returns a new State at the given prefix. The
    32  // following files are stored:
    33  // 	- {prefix}.json: the current state
    34  // 	- {prefix}.lock: the POSIX lock file
    35  // 	- {prefix}.bak: the previous state
    36  func Open(prefix string) (*File, error) {
    37  	f := &File{prefix: prefix}
    38  	os.MkdirAll(filepath.Dir(prefix), 0777) // best-effort
    39  	var err error
    40  	f.lockfd, err = syscall.Open(prefix+".lock", syscall.O_CREAT|syscall.O_RDWR, 0777)
    41  	if err != nil {
    42  		return nil, err
    43  	}
    44  	return f, nil
    45  }
    46  
    47  // Marshal opens a State, marshals v into it, and then closes it.
    48  func Marshal(prefix string, v interface{}) error {
    49  	file, err := Open(prefix)
    50  	if err != nil {
    51  		return err
    52  	}
    53  	err = file.Marshal(v)
    54  	file.Close()
    55  	return err
    56  }
    57  
    58  // Unmarshal opens a State, unmarshals it into v, and then closes it.
    59  func Unmarshal(prefix string, v interface{}) error {
    60  	file, err := Open(prefix)
    61  	if err != nil {
    62  		return err
    63  	}
    64  	err = file.Unmarshal(v)
    65  	file.Close()
    66  	return err
    67  }
    68  
    69  // Lock locks the state, both inside of the process and outside. Lock
    70  // relies on POSIX flock, which may not be available on all
    71  // filesystems, notably NFS and SMB.
    72  func (f *File) Lock() error {
    73  	f.mu.Lock()
    74  	if err := syscall.Flock(f.lockfd, syscall.LOCK_EX); err != nil {
    75  		f.mu.Unlock()
    76  		return err
    77  	}
    78  	return nil
    79  }
    80  
    81  // Unlock unlocks the state.
    82  func (f *File) Unlock() error {
    83  	if err := syscall.Flock(f.lockfd, syscall.LOCK_UN); err != nil {
    84  		return err
    85  	}
    86  	f.mu.Unlock()
    87  	return nil
    88  }
    89  
    90  // LockLocal locks local access to state.
    91  func (f *File) LockLocal() {
    92  	f.mu.Lock()
    93  }
    94  
    95  // UnlockLocal unlocks local access to state.
    96  func (f *File) UnlockLocal() {
    97  	f.mu.Unlock()
    98  }
    99  
   100  // Marshal atomically stores the JSON-encoded representation of v to
   101  // the current state. It is only stored when Marshal returns a nil
   102  // error.
   103  func (f *File) Marshal(v interface{}) error {
   104  	w, err := ioutil.TempFile(filepath.Dir(f.prefix), filepath.Base(f.prefix)+".write")
   105  	if err != nil {
   106  		return err
   107  	}
   108  	if err := json.NewEncoder(w).Encode(v); err != nil {
   109  		w.Close()
   110  		return err
   111  	}
   112  	if err := w.Close(); err != nil {
   113  		return err
   114  	}
   115  	os.Remove(f.prefix + ".bak")
   116  	os.Link(f.prefix+".json", f.prefix+".bak")
   117  	return os.Rename(w.Name(), f.prefix+".json")
   118  }
   119  
   120  // Unmarshal decodes the current state into v. Unmarshal returns
   121  // ErrNoState if no state is stored.
   122  func (f *File) Unmarshal(v interface{}) error {
   123  	w, err := os.Open(f.prefix + ".json")
   124  	if os.IsNotExist(err) {
   125  		return ErrNoState
   126  	} else if err != nil {
   127  		return err
   128  	}
   129  	defer w.Close()
   130  	return json.NewDecoder(w).Decode(v)
   131  }
   132  
   133  // Close releases resources associated with this State instance.
   134  func (f *File) Close() error {
   135  	return syscall.Close(f.lockfd)
   136  }