go.mondoo.com/cnquery@v0.0.0-20231005093811-59568235f6ea/mql/internal/waitgroup_test.go (about)

     1  // Copyright (c) Mondoo, Inc.
     2  // SPDX-License-Identifier: BUSL-1.1
     3  
     4  package internal
     5  
     6  import (
     7  	"sync"
     8  	"testing"
     9  
    10  	"github.com/stretchr/testify/require"
    11  )
    12  
    13  func TestWaitGroupInvalidUsage(t *testing.T) {
    14  	t.Run("calling Add with an active workID panics", func(t *testing.T) {
    15  		wg := NewWaitGroup()
    16  		wg.Add("foo")
    17  		require.Panics(t, func() { wg.Add("foo") })
    18  	})
    19  
    20  	t.Run("calling Done with an id that was never added panics", func(t *testing.T) {
    21  		wg := NewWaitGroup()
    22  		require.Panics(t, func() { wg.Done("foo") })
    23  	})
    24  }
    25  
    26  func TestWaitGroup(t *testing.T) {
    27  	t.Run("finishing completed workIDs unblocks Wait", func(t *testing.T) {
    28  		signalGoroutineStarted := &sync.WaitGroup{}
    29  		signalFinished := &sync.WaitGroup{}
    30  		wg := NewWaitGroup()
    31  		wg.Add("foo")
    32  
    33  		signalFinished.Add(1)
    34  		signalGoroutineStarted.Add(1)
    35  		go func() {
    36  			signalGoroutineStarted.Done()
    37  			wg.Wait()
    38  			signalFinished.Done()
    39  		}()
    40  
    41  		wg.Done("foo")
    42  		signalFinished.Wait()
    43  
    44  		stats := wg.Stats()
    45  		require.Equal(t, WaitGroupStats{
    46  			NumAdded:  1,
    47  			NumActive: 0,
    48  			NumDone:   1,
    49  		}, stats)
    50  
    51  		require.Equal(t, wg.IsDecommissioned(), false)
    52  	})
    53  
    54  	t.Run("decommissioning unblocks Wait", func(t *testing.T) {
    55  		signalGoroutineStarted := &sync.WaitGroup{}
    56  		signalFinished := &sync.WaitGroup{}
    57  		wg := NewWaitGroup()
    58  		wg.Add("foo")
    59  		wg.Add("bar")
    60  
    61  		signalFinished.Add(1)
    62  		signalGoroutineStarted.Add(1)
    63  		go func() {
    64  			signalGoroutineStarted.Done()
    65  			wg.Wait()
    66  			signalFinished.Done()
    67  		}()
    68  
    69  		signalGoroutineStarted.Wait()
    70  		wg.Decommission()
    71  		signalFinished.Wait()
    72  
    73  		stats := wg.Stats()
    74  		require.Equal(t, WaitGroupStats{
    75  			NumAdded:  2,
    76  			NumActive: 2,
    77  			NumDone:   0,
    78  		}, stats)
    79  
    80  		require.Equal(t, wg.IsDecommissioned(), true)
    81  	})
    82  
    83  	t.Run("usable after decommission", func(t *testing.T) {
    84  		// We want to make sure you can still done things even
    85  		// after decommissioning. The wait group should still
    86  		// never block on Wait after being decommissioned.
    87  
    88  		wg := NewWaitGroup()
    89  		wg.Add("foo")
    90  		wg.Add("bar")
    91  
    92  		wg.Decommission()
    93  		wg.Wait() // Does not block
    94  
    95  		wg.Done("foo")
    96  		wg.Add("baz")
    97  		wg.Wait()
    98  
    99  		stats := wg.Stats()
   100  		require.Equal(t, WaitGroupStats{
   101  			NumAdded:  3,
   102  			NumActive: 2,
   103  			NumDone:   1,
   104  		}, stats)
   105  
   106  		require.Equal(t, wg.IsDecommissioned(), true)
   107  	})
   108  }