github.com/grailbio/bigslice@v0.0.0-20230519005545-30c4c12152ad/exec/task_test.go (about)

     1  // Copyright 2019 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache 2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  package exec
     6  
     7  import (
     8  	"math/rand"
     9  	"reflect"
    10  	"sync"
    11  	"testing"
    12  )
    13  
    14  // TestTaskSubscriber verifies that task subscribers receive all tasks whose
    15  // state changes.
    16  func TestTaskSubscriber(t *testing.T) {
    17  	const (
    18  		numTasks   = 100000
    19  		numWriters = 8
    20  	)
    21  	var (
    22  		sub   = NewTaskSubscriber()
    23  		unsub = NewTaskSubscriber()
    24  		tasks = make([]*Task, numTasks)
    25  	)
    26  	for i := range tasks {
    27  		tasks[i] = &Task{}
    28  		tasks[i].Subscribe(sub)
    29  		// Throw in a subscriber that is immediately unsubscribed to make sure
    30  		// it doesn't gum up the works.
    31  		tasks[i].Subscribe(unsub)
    32  		tasks[i].Unsubscribe(unsub)
    33  	}
    34  	var (
    35  		mu      sync.Mutex
    36  		want    = make(map[*Task]bool)
    37  		writeWG sync.WaitGroup
    38  	)
    39  	for i := 0; i < numWriters; i++ {
    40  		writeWG.Add(1)
    41  		go func() {
    42  			defer writeWG.Done()
    43  			for j := 0; j < numTasks/numWriters/2; j++ {
    44  				task := tasks[rand.Intn(len(tasks))]
    45  				newState := TaskState(1 + rand.Intn(int(maxState)-1))
    46  				task.Set(newState)
    47  				mu.Lock()
    48  				want[task] = true
    49  				mu.Unlock()
    50  			}
    51  		}()
    52  	}
    53  	var (
    54  		got    = make(map[*Task]bool)
    55  		donec  = make(chan struct{})
    56  		readWG sync.WaitGroup
    57  	)
    58  	readWG.Add(1)
    59  	go func() {
    60  		defer readWG.Done()
    61  		for {
    62  			select {
    63  			case <-sub.Ready():
    64  				for _, task := range sub.Tasks() {
    65  					got[task] = true
    66  				}
    67  			case <-donec:
    68  				// Drain.
    69  				for {
    70  					select {
    71  					case <-sub.Ready():
    72  						for _, task := range sub.Tasks() {
    73  							got[task] = true
    74  						}
    75  					default:
    76  						return
    77  					}
    78  				}
    79  			}
    80  		}
    81  	}()
    82  	writeWG.Wait()
    83  	close(donec)
    84  	readWG.Wait()
    85  	if !reflect.DeepEqual(got, want) {
    86  		t.Logf("len(got), len(want): %d, %d", len(got), len(want))
    87  		t.Errorf("modified task was not seen by subscriber")
    88  	}
    89  	// The unsubscribed subscriber should see nothing.
    90  	if got, want := len(unsub.Tasks()), 0; got != want {
    91  		t.Errorf("got %v, want %v", got, want)
    92  	}
    93  }