github.com/cloudwego/kitex@v0.9.0/pkg/gofunc/go.go (about)

     1  /*
     2   * Copyright 2021 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package gofunc
    18  
    19  import (
    20  	"context"
    21  	"runtime/debug"
    22  	"sync"
    23  
    24  	"github.com/bytedance/gopkg/util/gopool"
    25  
    26  	"github.com/cloudwego/kitex/pkg/klog"
    27  	"github.com/cloudwego/kitex/pkg/profiler"
    28  )
    29  
    30  // GoTask is used to spawn a new task.
    31  type GoTask func(context.Context, func())
    32  
    33  // GoFunc is the default func used globally.
    34  var GoFunc GoTask
    35  
    36  func init() {
    37  	GoFunc = func(ctx context.Context, f func()) {
    38  		gopool.CtxGo(ctx, func() {
    39  			profiler.Tag(ctx)
    40  			f()
    41  			profiler.Untag(ctx)
    42  		})
    43  	}
    44  }
    45  
    46  // RecoverGoFuncWithInfo is the go func with recover panic and service info.
    47  // It is used for panic defence and output key info for troubleshooting.
    48  func RecoverGoFuncWithInfo(ctx context.Context, task func(), info *Info) {
    49  	GoFunc(ctx, func() {
    50  		defer func() {
    51  			if panicErr := recover(); panicErr != nil {
    52  				if info.RemoteService == "" {
    53  					info.RemoteService = "unknown"
    54  				}
    55  				if info.RemoteAddr == "" {
    56  					info.RemoteAddr = "unknown"
    57  				}
    58  				stack := string(debug.Stack())
    59  				klog.CtxErrorf(ctx, "KITEX: panic happened, remoteService=%s remoteAddress=%s error=%v\nstack=%s",
    60  					info.RemoteService, info.RemoteAddr, panicErr, stack)
    61  
    62  				phLock.RLock()
    63  				if panicHandler != nil {
    64  					panicHandler(info, panicErr, stack)
    65  				}
    66  				phLock.RUnlock()
    67  			}
    68  			infoPool.Put(info)
    69  		}()
    70  
    71  		task()
    72  	})
    73  }
    74  
    75  // SetPanicHandler is used to do something when panic happen, for example do metric report.
    76  func SetPanicHandler(hdlr func(info *Info, panicErr interface{}, panicStack string)) {
    77  	phLock.Lock()
    78  	panicHandler = hdlr
    79  	phLock.Unlock()
    80  }
    81  
    82  // NewBasicInfo is to new Info with remoteService  and remoteAddr.
    83  func NewBasicInfo(remoteService, remoteAddr string) *Info {
    84  	info := infoPool.Get().(*Info)
    85  	info.RemoteService = remoteService
    86  	info.RemoteAddr = remoteAddr
    87  	return info
    88  }
    89  
    90  // Info is used to pass key info to go func, which is convenient for troubleshooting
    91  type Info struct {
    92  	RemoteService string
    93  	RemoteAddr    string
    94  }
    95  
    96  var (
    97  	// EmptyInfo is empty Info which is convenient to reuse
    98  	EmptyInfo = &Info{}
    99  
   100  	panicHandler func(info *Info, panicErr interface{}, panicStack string)
   101  	phLock       sync.RWMutex
   102  
   103  	infoPool = &sync.Pool{
   104  		New: func() interface{} {
   105  			return new(Info)
   106  		},
   107  	}
   108  )