github.com/mailru/activerecord@v1.12.2/pkg/iproto/syncutil/taskgroup_test.go (about)

     1  package syncutil
     2  
     3  import (
     4  	"sync/atomic"
     5  	"testing"
     6  	"time"
     7  
     8  	"golang.org/x/net/context"
     9  )
    10  
    11  var bg = context.Background()
    12  
    13  func TestTaskGroupDoDeadlock(t *testing.T) {
    14  	sem := make(chan struct{}, 1)
    15  	tg := TaskGroup{
    16  		N: 2,
    17  		Goer: GoerFn(func(ctx context.Context, task func()) error {
    18  			sem <- struct{}{}
    19  			go func() {
    20  				defer func() { <-sem }()
    21  				task()
    22  			}()
    23  			return nil
    24  		}),
    25  	}
    26  
    27  	var (
    28  		done = make(chan struct{})
    29  		task = make(chan int, 2)
    30  	)
    31  	go func() {
    32  		defer close(done)
    33  		tg.Do(bg, 2, func(ctx context.Context, i int) error {
    34  			task <- i
    35  			return nil
    36  		})
    37  	}()
    38  	select {
    39  	case <-done:
    40  		t.Errorf("Do() returned; want deadlock")
    41  	case <-time.After(time.Second):
    42  	}
    43  	if n := len(task); n != 1 {
    44  		t.Fatalf("want only one task to be executed; got %d", n)
    45  	}
    46  	if i := <-task; i != 0 {
    47  		t.Fatalf("want task #%d be executed; got #%d", 0, i)
    48  	}
    49  }
    50  
    51  func TestTaskGroupDo(t *testing.T) {
    52  	const N = 8
    53  	s := TaskGroup{
    54  		N: N,
    55  	}
    56  
    57  	sleep := make(chan struct{})
    58  	ret := s.Do(bg, N, func(ctx context.Context, i int) error {
    59  		<-sleep
    60  		return nil
    61  	})
    62  
    63  	for i := 0; i < 100; i++ {
    64  		time.Sleep(time.Millisecond)
    65  		s.Do(bg, N, func(_ context.Context, _ int) error {
    66  			panic("must not be called")
    67  		})
    68  	}
    69  
    70  	close(sleep)
    71  	if err := WaitPending(bg, ret); err != nil {
    72  		t.Fatalf("unexpected error: %v", err)
    73  	}
    74  }
    75  
    76  func TestTaskGroupCancel(t *testing.T) {
    77  	s := TaskGroup{
    78  		N: 1,
    79  	}
    80  
    81  	time.AfterFunc(50*time.Millisecond, func() {
    82  		s.Cancel()
    83  	})
    84  	a := DoOne(bg, &s, func(ctx context.Context) error {
    85  		<-ctx.Done()
    86  		return ctx.Err()
    87  	})
    88  	b := DoOne(bg, &s, func(ctx context.Context) error {
    89  		panic("must not be called")
    90  	})
    91  
    92  	if b != a {
    93  		t.Fatalf("unexpected exec")
    94  	}
    95  
    96  	if err, want := <-a, context.Canceled; err != want {
    97  		t.Errorf("got %v; want %v", err, want)
    98  	}
    99  
   100  	c := DoOne(bg, &s, func(ctx context.Context) error {
   101  		return nil
   102  	})
   103  
   104  	if err := <-c; err != nil {
   105  		t.Fatal(err)
   106  	}
   107  }
   108  
   109  func TestTaskGroupSplit(t *testing.T) {
   110  	s := TaskGroup{
   111  		N: 2,
   112  	}
   113  	var (
   114  		a = make(chan struct{}, 1)
   115  		b = make(chan struct{}, 1)
   116  		n = new(int32)
   117  	)
   118  	s.Do(bg, 2, func(ctx context.Context, i int) error {
   119  		atomic.AddInt32(n, 1)
   120  		switch i {
   121  		case 0:
   122  			<-a
   123  		case 1:
   124  			<-b
   125  		}
   126  		return nil
   127  	})
   128  
   129  	// Release first task.
   130  	a <- struct{}{}
   131  	time.Sleep(10 * time.Millisecond)
   132  	s.Do(bg, 2, func(ctx context.Context, i int) error {
   133  		atomic.AddInt32(n, 1)
   134  		if i != 0 {
   135  			t.Fatalf("unexpected index: %d", i)
   136  		}
   137  		<-a
   138  		return nil
   139  	})
   140  
   141  	// Release second task.
   142  	b <- struct{}{}
   143  	time.Sleep(10 * time.Millisecond)
   144  	s.Do(bg, 2, func(ctx context.Context, i int) error {
   145  		atomic.AddInt32(n, 1)
   146  		if i != 1 {
   147  			t.Fatalf("unexpected index: %d", i)
   148  		}
   149  		<-b
   150  		return nil
   151  	})
   152  
   153  	time.Sleep(10 * time.Millisecond)
   154  	if m := atomic.LoadInt32(n); m != 4 {
   155  		t.Fatalf("unexpected number of executed tasks: %d", m)
   156  	}
   157  }
   158  
   159  func DoOne(ctx context.Context, s *TaskGroup, cb func(context.Context) error) <-chan error {
   160  	ps := s.Do(ctx, 1, func(ctx context.Context, _ int) error {
   161  		return cb(ctx)
   162  	})
   163  	return ps[0]
   164  }
   165  
   166  func WaitPending(ctx context.Context, chs []<-chan error) error {
   167  	var fail error
   168  	for _, ch := range chs {
   169  		select {
   170  		case <-ctx.Done():
   171  			return ctx.Err()
   172  		case err := <-ch:
   173  			if err != nil && fail == nil {
   174  				fail = err
   175  			}
   176  		}
   177  	}
   178  	return fail
   179  }