github.com/pinpoint-apm/pinpoint-go-agent@v1.4.1-0.20240110120318-a50c2eb18c8c/command.go (about)

     1  package pinpoint
     2  
     3  import (
     4  	pb "github.com/pinpoint-apm/pinpoint-go-agent/protobuf"
     5  	"io"
     6  	"sync"
     7  	"sync/atomic"
     8  	"time"
     9  )
    10  
    11  var (
    12  	gAtcStreamCount    int32
    13  	realTimeActiveSpan sync.Map
    14  )
    15  
    16  type activeSpanInfo struct {
    17  	startTime  time.Time
    18  	txId       string
    19  	entryPoint string
    20  	sampled    bool
    21  }
    22  
    23  func (agent *agent) runCommandService() {
    24  	Log("cmd").Infof("start command goroutine")
    25  	defer agent.wg.Done()
    26  	gAtcStreamCount = 0
    27  
    28  	for agent.enable {
    29  		stream := agent.cmdGrpc.newCommandStreamWithRetry()
    30  		err := stream.sendCommandMessage()
    31  		if err != nil {
    32  			if err != io.EOF {
    33  				Log("cmd").Errorf("send command message - %v", err)
    34  			}
    35  			stream.close()
    36  			continue
    37  		}
    38  
    39  		for agent.enable {
    40  			cmdReq, err := stream.recvCommandRequest()
    41  			if err != nil {
    42  				if agent.enable && err != io.EOF {
    43  					Log("cmd").Warnf("recv command request - %v", err)
    44  				}
    45  				break
    46  			}
    47  
    48  			reqId := cmdReq.GetRequestId()
    49  			Log("cmd").Infof("command request: %v, %v", cmdReq, reqId)
    50  
    51  			switch cmdReq.Command.(type) {
    52  			case *pb.PCmdRequest_CommandEcho:
    53  				msg := cmdReq.GetCommandEcho().GetMessage()
    54  				agent.cmdGrpc.sendEcho(reqId, msg)
    55  				break
    56  			case *pb.PCmdRequest_CommandActiveThreadCount:
    57  				atcStream := agent.cmdGrpc.newActiveThreadCountStream(reqId)
    58  				go agent.sendActiveThreadCount(atcStream)
    59  				break
    60  			case *pb.PCmdRequest_CommandActiveThreadDump:
    61  				if c := cmdReq.GetCommandActiveThreadDump(); c != nil {
    62  					limit := c.GetLimit()
    63  					threadName := c.GetThreadName()
    64  					localId := c.GetLocalTraceId()
    65  					agent.cmdGrpc.sendActiveThreadDump(reqId, limit, threadName, localId, dumpGoroutine())
    66  				}
    67  				break
    68  			case *pb.PCmdRequest_CommandActiveThreadLightDump:
    69  				if c := cmdReq.GetCommandActiveThreadLightDump(); c != nil {
    70  					agent.cmdGrpc.sendActiveThreadLightDump(reqId, c.GetLimit(), dumpGoroutine())
    71  				}
    72  				break
    73  			case nil:
    74  			default:
    75  				break
    76  			}
    77  		}
    78  
    79  		stream.close()
    80  	}
    81  
    82  	Log("cmd").Infof("end command goroutine")
    83  }
    84  
    85  func (agent *agent) sendActiveThreadCount(s *activeThreadCountStream) {
    86  	atomic.AddInt32(&gAtcStreamCount, 1)
    87  	Log("cmd").Infof("active thread count stream goroutine start: %d, %d", s.reqId, gAtcStreamCount)
    88  
    89  	for agent.enable {
    90  		err := s.sendActiveThreadCount()
    91  		if err != nil {
    92  			if err != io.EOF {
    93  				Log("cmd").Errorf("send active thread count - %d, %v", s.reqId, err)
    94  			}
    95  			break
    96  		}
    97  		time.Sleep(1 * time.Second)
    98  	}
    99  	s.close()
   100  
   101  	atomic.AddInt32(&gAtcStreamCount, -1)
   102  	Log("cmd").Infof("active thread count stream goroutine finish: %d, %d", s.reqId, gAtcStreamCount)
   103  }
   104  
   105  func addRealTimeSampledActiveSpan(span *span) {
   106  	if gAtcStreamCount > 0 {
   107  		span.goroutineId = curGoroutineID()
   108  		s := &activeSpanInfo{span.startTime, span.txId.String(), span.rpcName, true}
   109  		realTimeActiveSpan.Store(span.goroutineId, s)
   110  	}
   111  }
   112  
   113  func dropRealTimeSampledActiveSpan(span *span) {
   114  	realTimeActiveSpan.Delete(span.goroutineId)
   115  }
   116  
   117  func addRealTimeUnSampledActiveSpan(span *noopSpan) {
   118  	if gAtcStreamCount > 0 {
   119  		span.goroutineId = curGoroutineID()
   120  		s := &activeSpanInfo{span.startTime, "", span.rpcName, false}
   121  		realTimeActiveSpan.Store(span.goroutineId, s)
   122  	}
   123  }
   124  
   125  func dropRealTimeUnSampledActiveSpan(span *noopSpan) {
   126  	realTimeActiveSpan.Delete(span.goroutineId)
   127  }
   128  
   129  func getActiveSpanCount(now time.Time) []int32 {
   130  	counts := []int32{0, 0, 0, 0}
   131  	realTimeActiveSpan.Range(func(k, v interface{}) bool {
   132  		s := v.(*activeSpanInfo)
   133  		d := now.Sub(s.startTime).Seconds()
   134  
   135  		if d < 1 {
   136  			counts[0]++
   137  		} else if d < 3 {
   138  			counts[1]++
   139  		} else if d < 5 {
   140  			counts[2]++
   141  		} else {
   142  			counts[3]++
   143  		}
   144  		return true
   145  	})
   146  
   147  	return counts
   148  }