github.com/Jeffail/benthos/v3@v3.65.0/lib/util/closable_pool.go (about)

     1  package util
     2  
     3  import (
     4  	"sort"
     5  	"time"
     6  
     7  	"github.com/Jeffail/benthos/v3/lib/types"
     8  )
     9  
    10  //------------------------------------------------------------------------------
    11  
    12  // ClosablePool keeps a reference to a pool of closable types and closes them in
    13  // tiers.
    14  type ClosablePool struct {
    15  	closables map[int][]types.Closable
    16  }
    17  
    18  // NewClosablePool creates a fresh pool of closable types.
    19  func NewClosablePool() *ClosablePool {
    20  	return &ClosablePool{
    21  		closables: make(map[int][]types.Closable),
    22  	}
    23  }
    24  
    25  //------------------------------------------------------------------------------
    26  
    27  // Add adds a closable type to the pool, tiers are used to partition and order
    28  // the closing of types (starting at the lowest tier and working upwards).
    29  // Closable types in a single tier are closed in the order that they are added.
    30  func (c *ClosablePool) Add(tier int, closable types.Closable) {
    31  	tierArray := []types.Closable{}
    32  	if t, ok := c.closables[tier]; ok {
    33  		tierArray = t
    34  	}
    35  	c.closables[tier] = append(tierArray, closable)
    36  }
    37  
    38  // Close attempts to close and clean up all stored closables in the determined
    39  // order. If the timeout is met whilst working through the pool there is no
    40  // indication of how far the pool has been progressed, thus this timeout should
    41  // only be used for preventing severe blocking of a service.
    42  func (c *ClosablePool) Close(timeout time.Duration) error {
    43  	started := time.Now()
    44  
    45  	tiers := []int{}
    46  	for i := range c.closables {
    47  		tiers = append(tiers, i)
    48  	}
    49  	sort.Ints(tiers)
    50  
    51  	for _, i := range tiers {
    52  		tier := c.closables[i]
    53  		for j := range tier {
    54  			tier[j].CloseAsync()
    55  		}
    56  		for j := range tier {
    57  			if err := tier[j].WaitForClose(timeout - time.Since(started)); err != nil {
    58  				return err
    59  			}
    60  		}
    61  		delete(c.closables, i)
    62  	}
    63  	return nil
    64  }
    65  
    66  //------------------------------------------------------------------------------