github.com/niedbalski/juju@v0.0.0-20190215020005-8ff100488e47/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 and associated *state.StatePool.
    18  type StateTracker interface {
    19  	// Use returns the wrapped StatePool, recording the use of
    20  	// it. ErrStateClosed is returned if the StatePool is closed.
    21  	Use() (*state.StatePool, error)
    22  
    23  	// Done records that there's one less user of the wrapped StatePool,
    24  	// closing it if there's no more users. ErrStateClosed is returned
    25  	// if the StatePool has already been closed (indicating that Done has
    26  	// called too many times).
    27  	Done() error
    28  
    29  	// Report is used to give details about what is going on with this state tracker.
    30  	Report() map[string]interface{}
    31  }
    32  
    33  // stateTracker wraps a *state.State, keeping a reference count and
    34  // closing the State and associated *state.StatePool when there are
    35  // no longer any references. It implements StateTracker.
    36  //
    37  // The reference count starts at 1. Done should be called exactly 1 +
    38  // number of calls to Use.
    39  type stateTracker struct {
    40  	mu         sync.Mutex
    41  	pool       *state.StatePool
    42  	references int
    43  }
    44  
    45  func newStateTracker(pool *state.StatePool) StateTracker {
    46  	return &stateTracker{
    47  		pool:       pool,
    48  		references: 1,
    49  	}
    50  }
    51  
    52  // Use implements StateTracker.
    53  func (c *stateTracker) Use() (*state.StatePool, error) {
    54  	c.mu.Lock()
    55  	defer c.mu.Unlock()
    56  
    57  	if c.references == 0 {
    58  		return nil, ErrStateClosed
    59  	}
    60  	c.references++
    61  	return c.pool, nil
    62  }
    63  
    64  // Done implements StateTracker.
    65  func (c *stateTracker) Done() error {
    66  	c.mu.Lock()
    67  	defer c.mu.Unlock()
    68  
    69  	if c.references == 0 {
    70  		return ErrStateClosed
    71  	}
    72  	c.references--
    73  	if c.references == 0 {
    74  		if err := c.pool.Close(); err != nil {
    75  			logger.Errorf("error when closing state pool: %v", err)
    76  		}
    77  	}
    78  	return nil
    79  }
    80  
    81  func (c *stateTracker) Report() map[string]interface{} {
    82  	c.mu.Lock()
    83  	defer c.mu.Unlock()
    84  	return c.pool.Report()
    85  }