github.com/niedbalski/juju@v0.0.0-20190215020005-8ff100488e47/worker/common/cleanup.go (about)

     1  // Copyright 2018 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package common
     5  
     6  import (
     7  	"sync"
     8  
     9  	"gopkg.in/juju/worker.v1"
    10  )
    11  
    12  // NewCleanupWorker returns a worker that ensures a cleanup function
    13  // is run after the underlying worker is finished.
    14  func NewCleanupWorker(w worker.Worker, cleanup func()) worker.Worker {
    15  	return &CleanupWorker{
    16  		Worker:  w,
    17  		cleanup: cleanup,
    18  	}
    19  }
    20  
    21  // CleanupWorker wraps another worker to ensure a func is run when it
    22  // is finished. (Public for manifolds that need access to the
    23  // wrapped worker for output.)
    24  type CleanupWorker struct {
    25  	worker.Worker
    26  	cleanupOnce sync.Once
    27  	cleanup     func()
    28  }
    29  
    30  // Wait ensures the cleanup func is run after the worker finishes.
    31  func (w *CleanupWorker) Wait() error {
    32  	err := w.Worker.Wait()
    33  	w.cleanupOnce.Do(w.cleanup)
    34  	return err
    35  }
    36  
    37  // Report implements dependency.Reporter.
    38  func (w *CleanupWorker) Report() map[string]interface{} {
    39  	if r, ok := w.Worker.(worker.Reporter); ok {
    40  		return r.Report()
    41  	}
    42  	return nil
    43  }