github.com/smallnest/goroutine@v1.1.1/gid_test.go (about)

     1  // Copyright ©2020 Dan Kortschak. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package goroutine
     6  
     7  import (
     8  	"fmt"
     9  	"runtime"
    10  	"strconv"
    11  	"strings"
    12  	"sync"
    13  	"testing"
    14  )
    15  
    16  func TestID(t *testing.T) {
    17  	got := ID()
    18  	want := goid()
    19  	if got != want {
    20  		t.Fatalf("unexpected id for main goroutine: got:%d want:%d", got, want)
    21  	}
    22  	var wg sync.WaitGroup
    23  	for i := 0; i < 1000000; i++ {
    24  		i := i
    25  		wg.Add(1)
    26  		go func() {
    27  			defer wg.Done()
    28  			got := ID()
    29  			want := goid()
    30  			if got != want {
    31  				t.Errorf("unexpected id for goroutine number %d: got:%d want:%d", i, got, want)
    32  			}
    33  		}()
    34  	}
    35  	wg.Wait()
    36  }
    37  
    38  // goid returns the goroutine ID extracted from a stack trace.
    39  func goid() int64 {
    40  	var buf [64]byte
    41  	n := runtime.Stack(buf[:], false)
    42  	idField := strings.Fields(strings.TrimPrefix(string(buf[:n]), "goroutine "))[0]
    43  	id, err := strconv.ParseInt(idField, 10, 64)
    44  	if err != nil {
    45  		panic(fmt.Sprintf("cannot get goroutine id: %v", err))
    46  	}
    47  	return id
    48  }
    49  
    50  func TestPID(t *testing.T) {
    51  	n := runtime.GOMAXPROCS(0)
    52  	got := PID()
    53  	if got < 0 || got >= n {
    54  		t.Fatalf("unexpected pid for main goroutine: got:%d want:0 < pid < %d", got, n)
    55  	}
    56  	var wg sync.WaitGroup
    57  	for i := 0; i < 1000000; i++ {
    58  		i := i
    59  		wg.Add(1)
    60  		go func() {
    61  			defer wg.Done()
    62  			got := PID()
    63  			if got < 0 || got >= n {
    64  				t.Fatalf("unexpected pid for goroutine #%d: got:%d want:0 <= pid < %d", i, got, n)
    65  			}
    66  		}()
    67  	}
    68  	wg.Wait()
    69  }