github.com/madokast/nopreempt@v1.0.3-0.20240125081507-3fa05aef38ed/nopreempt_test.go (about)

     1  package nopreempt
     2  
     3  import (
     4  	"bytes"
     5  	"math/rand"
     6  	"runtime"
     7  	"strconv"
     8  	"sync"
     9  	"testing"
    10  	"time"
    11  )
    12  
    13  func TestGoId(t *testing.T) {
    14  	var wg sync.WaitGroup
    15  	for i := 0; i < 256; i++ {
    16  		wg.Add(1)
    17  		go func() {
    18  			defer wg.Done()
    19  			// t.Log("slow", goid())
    20  			// t.Log("quick", GetGID())
    21  			if goid() != GetGID() {
    22  				t.Fail()
    23  			}
    24  		}()
    25  	}
    26  	wg.Wait()
    27  }
    28  
    29  func TestMId(t *testing.T) {
    30  	const mp = 1
    31  	omp := runtime.GOMAXPROCS(mp)
    32  	defer runtime.GOMAXPROCS(omp)
    33  	set := map[int64]struct{}{}
    34  	var wg sync.WaitGroup
    35  	for i := 0; i < 3; i++ {
    36  		wg.Add(1)
    37  		go func() {
    38  			defer wg.Done()
    39  			set[GetMID()] = struct{}{}
    40  		}()
    41  	}
    42  	wg.Wait()
    43  	t.Log(set)
    44  	if len(set) != mp {
    45  		t.Fail()
    46  	}
    47  }
    48  
    49  func TestMId2(t *testing.T) {
    50  	const mp = 2
    51  	omp := runtime.GOMAXPROCS(mp)
    52  	defer runtime.GOMAXPROCS(omp)
    53  	set := map[int64]struct{}{}
    54  	var wg sync.WaitGroup
    55  	var mu sync.Mutex
    56  	for i := 0; i < 3; i++ {
    57  		wg.Add(1)
    58  		go func() {
    59  			defer wg.Done()
    60  			mu.Lock()
    61  			defer mu.Unlock()
    62  			set[GetMID()] = struct{}{}
    63  		}()
    64  	}
    65  	wg.Wait()
    66  	t.Log(set)
    67  	if len(set) > mp {
    68  		t.Fail()
    69  	}
    70  }
    71  
    72  func TestMId3(t *testing.T) {
    73  	const mp = 3
    74  	omp := runtime.GOMAXPROCS(mp)
    75  	defer runtime.GOMAXPROCS(omp)
    76  	t.Log("old GOMAXPROCS", omp)
    77  	set := map[int64]struct{}{}
    78  	var wg sync.WaitGroup
    79  	var mu sync.Mutex
    80  	for i := 0; i < omp*3; i++ {
    81  		wg.Add(1)
    82  		go func() {
    83  			defer wg.Done()
    84  			mu.Lock()
    85  			set[GetMID()] = struct{}{}
    86  			mu.Unlock()
    87  
    88  			var s float64
    89  			for k := 0; k < 1000000; k++ {
    90  				s *= rand.Float64()
    91  			}
    92  		}()
    93  	}
    94  	wg.Wait()
    95  	t.Log(set)
    96  	if len(set) > mp {
    97  		t.Fail()
    98  	}
    99  }
   100  
   101  func TestMId10(t *testing.T) {
   102  	const mp = 10
   103  	omp := runtime.GOMAXPROCS(mp)
   104  	defer runtime.GOMAXPROCS(omp)
   105  	t.Log("old GOMAXPROCS", omp)
   106  	set := map[int64]struct{}{}
   107  	var wg sync.WaitGroup
   108  	var mu sync.Mutex
   109  	for i := 0; i < omp*3; i++ {
   110  		wg.Add(1)
   111  		go func() {
   112  			defer wg.Done()
   113  			mu.Lock()
   114  			set[GetMID()] = struct{}{}
   115  			mu.Unlock()
   116  
   117  			var s float64
   118  			for k := 0; k < 1000000; k++ {
   119  				s *= rand.Float64()
   120  			}
   121  		}()
   122  	}
   123  	wg.Wait()
   124  	t.Log(set)
   125  	if len(set) > mp {
   126  		t.Fail()
   127  	}
   128  }
   129  
   130  func TestNoPreempt(t *testing.T) {
   131  	omp := runtime.GOMAXPROCS(1)
   132  	defer runtime.GOMAXPROCS(omp)
   133  
   134  	var wg sync.WaitGroup
   135  	wg.Add(1)
   136  	start := time.Now()
   137  	go func() {
   138  		for time.Since(start) < 2*time.Second {
   139  			t.Log("busy")
   140  			var s float64
   141  			for k := 0; k < 100000000; k++ {
   142  				s *= rand.Float64()
   143  			}
   144  		}
   145  		wg.Done()
   146  	}()
   147  	time.Sleep(100 * time.Millisecond)
   148  	t.Log("exit")
   149  	if time.Since(start) < time.Second {
   150  		t.Log("preempted")
   151  	} else {
   152  		t.Error("cannot preempt")
   153  		t.Fail()
   154  	}
   155  	wg.Wait()
   156  }
   157  
   158  func TestPreempt(t *testing.T) {
   159  	omp := runtime.GOMAXPROCS(1)
   160  	defer runtime.GOMAXPROCS(omp)
   161  
   162  	var wg sync.WaitGroup
   163  	wg.Add(1)
   164  	start := time.Now()
   165  	go func() {
   166  		mp := AcquireM()
   167  		defer mp.Release()
   168  		for time.Since(start) < 2*time.Second {
   169  			t.Log("busy")
   170  			var s float64
   171  			for k := 0; k < 100000000; k++ {
   172  				s *= rand.Float64()
   173  			}
   174  		}
   175  		wg.Done()
   176  	}()
   177  	time.Sleep(100 * time.Millisecond)
   178  	t.Log("exit")
   179  	if time.Since(start) < 2*time.Second {
   180  		t.Error("preempted")
   181  		t.Fail()
   182  	} else {
   183  		t.Log("cannot preempt")
   184  	}
   185  	wg.Wait()
   186  }
   187  
   188  func goid() int64 {
   189  	buf := make([]byte, len("goroutine ddddddddd"))
   190  	runtime.Stack(buf, false)
   191  	buf = buf[len("goroutine "):]
   192  	buf = buf[:bytes.IndexByte(buf, ' ')]
   193  	gid, _ := strconv.ParseInt(string(buf), 10, 64)
   194  	return gid
   195  }