github.com/wallyworld/juju@v0.0.0-20161013125918-6cf1bc9d917a/worker/state/statetracker.go (about)

     1  // Copyright 2016 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package state
     5  
     6  import (
     7  	"sync"
     8  
     9  	"github.com/juju/errors"
    10  
    11  	"github.com/juju/juju/state"
    12  )
    13  
    14  var ErrStateClosed = errors.New("state closed")
    15  
    16  // StateTracker describes a type which wraps and manages the lifetime
    17  // of a *state.State.
    18  type StateTracker interface {
    19  	// Use returns wrapped State, recording the use of
    20  	// it. ErrStateClosed is returned if the State is closed.
    21  	Use() (*state.State, error)
    22  
    23  	// Done records that there's one less user of the wrapped State,
    24  	// closing it if there's no more users. ErrStateClosed is returned
    25  	// if the State has already been closed (indicating that Done has
    26  	// called too many times).
    27  	Done() error
    28  }
    29  
    30  // stateTracker wraps a *state.State, keeping a reference count and
    31  // closing the State when there are no longer any references to it. It
    32  // implements StateTracker.
    33  //
    34  // The reference count starts at 1. Done should be called exactly 1 +
    35  // number of calls to Use.
    36  type stateTracker struct {
    37  	mu         sync.Mutex
    38  	st         *state.State
    39  	references int
    40  }
    41  
    42  func newStateTracker(st *state.State) StateTracker {
    43  	return &stateTracker{
    44  		st:         st,
    45  		references: 1,
    46  	}
    47  }
    48  
    49  // Use implements StateTracker.
    50  func (c *stateTracker) Use() (*state.State, error) {
    51  	c.mu.Lock()
    52  	defer c.mu.Unlock()
    53  
    54  	if c.references == 0 {
    55  		return nil, ErrStateClosed
    56  	}
    57  	c.references++
    58  	return c.st, nil
    59  }
    60  
    61  // Done implements StateTracker.
    62  func (c *stateTracker) Done() error {
    63  	c.mu.Lock()
    64  	defer c.mu.Unlock()
    65  
    66  	if c.references == 0 {
    67  		return ErrStateClosed
    68  	}
    69  	c.references--
    70  	if c.references == 0 {
    71  		err := c.st.Close()
    72  		if err != nil {
    73  			logger.Errorf("error when closing state: %v", err)
    74  		}
    75  	}
    76  	return nil
    77  }