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

     1  package util
     2  
     3  import (
     4  	"testing"
     5  	"time"
     6  
     7  	"github.com/Jeffail/benthos/v3/lib/types"
     8  )
     9  
    10  //------------------------------------------------------------------------------
    11  
    12  type closable struct {
    13  	globalCloseCounter *int
    14  	globalWaitCounter  *int
    15  
    16  	waitedAt int
    17  	closedAt int
    18  
    19  	waitFor time.Duration
    20  }
    21  
    22  func (c *closable) CloseAsync() {
    23  	c.waitedAt = *c.globalWaitCounter
    24  	*c.globalWaitCounter++
    25  }
    26  
    27  func (c *closable) WaitForClose(tout time.Duration) error {
    28  	if c.waitFor > 0 {
    29  		if tout < c.waitFor {
    30  			<-time.After(tout)
    31  			return types.ErrTimeout
    32  		}
    33  		<-time.After(c.waitFor)
    34  	}
    35  	c.closedAt = *c.globalCloseCounter
    36  	*c.globalCloseCounter++
    37  	return nil
    38  }
    39  
    40  //------------------------------------------------------------------------------
    41  
    42  func TestClosablePoolOrdering(t *testing.T) {
    43  	n := 100
    44  
    45  	closables := []*closable{}
    46  
    47  	pool := NewClosablePool()
    48  
    49  	closeCount, waitCount := 0, 0
    50  
    51  	for i := 0; i < n; i++ {
    52  		closables = append(closables, &closable{
    53  			globalCloseCounter: &closeCount,
    54  			globalWaitCounter:  &waitCount,
    55  		})
    56  	}
    57  
    58  	for i := range closables {
    59  		pool.Add(0, closables[i])
    60  	}
    61  
    62  	pool.Close(time.Second)
    63  
    64  	if waitCount != n {
    65  		t.Errorf("Wrong global wait count: %v != %v", waitCount, n)
    66  	}
    67  	if closeCount != n {
    68  		t.Errorf("Wrong global close count: %v != %v", closeCount, n)
    69  	}
    70  
    71  	for i := range closables {
    72  		if actual := closables[i].waitedAt; actual != i {
    73  			t.Errorf("Wrong closable wait index: %v != %v", actual, i)
    74  		}
    75  		if actual := closables[i].closedAt; actual != i {
    76  			t.Errorf("Wrong closable closed index: %v != %v", actual, i)
    77  		}
    78  	}
    79  }
    80  
    81  func TestClosablePoolTierOrdering(t *testing.T) {
    82  	n, tiers, closeCount, waitCount := 100, 5, 0, 0
    83  
    84  	closables := [][]*closable{}
    85  	m := map[int]struct{}{}
    86  
    87  	pool := NewClosablePool()
    88  
    89  	for i := 0; i < tiers; i++ {
    90  		tClosables := []*closable{}
    91  		m[i] = struct{}{}
    92  		for j := 0; j < n; j++ {
    93  			tClosables = append(tClosables, &closable{
    94  				globalCloseCounter: &closeCount,
    95  				globalWaitCounter:  &waitCount,
    96  			})
    97  		}
    98  		closables = append(closables, tClosables)
    99  	}
   100  
   101  	// Random iteration
   102  	for i := range m {
   103  		tClosables := closables[i]
   104  		for j := range tClosables {
   105  			pool.Add(i, tClosables[j])
   106  		}
   107  	}
   108  
   109  	pool.Close(time.Second)
   110  
   111  	if waitCount != n*tiers {
   112  		t.Errorf("Wrong global wait count: %v != %v", waitCount, n*tiers)
   113  	}
   114  	if closeCount != n*tiers {
   115  		t.Errorf("Wrong global close count: %v != %v", closeCount, n*tiers)
   116  	}
   117  
   118  	for i := range closables {
   119  		for j := range closables[i] {
   120  			if actual := closables[i][j].waitedAt; actual != i*n+j {
   121  				t.Errorf("Wrong closable wait index: %v != %v", actual, i*n+j)
   122  			}
   123  			if actual := closables[i][j].closedAt; actual != i*n+j {
   124  				t.Errorf("Wrong closable closed index: %v != %v", actual, i*n+j)
   125  			}
   126  		}
   127  	}
   128  }
   129  
   130  //------------------------------------------------------------------------------