github.com/cloudwego/hertz@v0.9.3/pkg/common/stackless/func.go (about)

     1  /*
     2   * Copyright 2022 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   * The MIT License (MIT)
    17   *
    18   * Copyright (c) 2015-present Aliaksandr Valialkin, VertaMedia, Kirill Danshin, Erik Dubbelboer, FastHTTP Authors
    19   *
    20   * Permission is hereby granted, free of charge, to any person obtaining a copy
    21   * of this software and associated documentation files (the "Software"), to deal
    22   * in the Software without restriction, including without limitation the rights
    23   * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
    24   * copies of the Software, and to permit persons to whom the Software is
    25   * furnished to do so, subject to the following conditions:
    26   *
    27   * The above copyright notice and this permission notice shall be included in
    28   * all copies or substantial portions of the Software.
    29   *
    30   * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    31   * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    32   * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    33   * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    34   * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    35   * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    36   * THE SOFTWARE.
    37   *
    38   * This file may have been modified by CloudWeGo authors. All CloudWeGo
    39   * Modifications are Copyright 2022 CloudWeGo Authors.
    40   */
    41  
    42  package stackless
    43  
    44  import (
    45  	"runtime"
    46  	"sync"
    47  )
    48  
    49  // NewFunc returns stackless wrapper for the function f.
    50  //
    51  // Unlike f, the returned stackless wrapper doesn't use stack space
    52  // on the goroutine that calls it.
    53  // The wrapper may save a lot of stack space if the following conditions
    54  // are met:
    55  //
    56  //   - f doesn't contain blocking calls on network, I/O or channels;
    57  //   - f uses a lot of stack space;
    58  //   - the wrapper is called from high number of concurrent goroutines.
    59  //
    60  // The stackless wrapper returns false if the call cannot be processed
    61  // at the moment due to high load.
    62  func NewFunc(f func(ctx interface{})) func(ctx interface{}) bool {
    63  	if f == nil {
    64  		panic("BUG: f cannot be nil")
    65  	}
    66  
    67  	funcWorkCh := make(chan *funcWork, runtime.GOMAXPROCS(-1)*2048)
    68  	onceInit := func() {
    69  		n := runtime.GOMAXPROCS(-1)
    70  		for i := 0; i < n; i++ {
    71  			go funcWorker(funcWorkCh, f)
    72  		}
    73  	}
    74  	var once sync.Once
    75  
    76  	return func(ctx interface{}) bool {
    77  		once.Do(onceInit)
    78  		fw := getFuncWork()
    79  		fw.ctx = ctx
    80  
    81  		select {
    82  		case funcWorkCh <- fw:
    83  		default:
    84  			putFuncWork(fw)
    85  			return false
    86  		}
    87  		<-fw.done
    88  		putFuncWork(fw)
    89  		return true
    90  	}
    91  }
    92  
    93  func funcWorker(funcWorkCh <-chan *funcWork, f func(ctx interface{})) {
    94  	for fw := range funcWorkCh {
    95  		f(fw.ctx)
    96  		fw.done <- struct{}{}
    97  	}
    98  }
    99  
   100  func getFuncWork() *funcWork {
   101  	v := funcWorkPool.Get()
   102  	if v == nil {
   103  		v = &funcWork{
   104  			done: make(chan struct{}, 1),
   105  		}
   106  	}
   107  	return v.(*funcWork)
   108  }
   109  
   110  func putFuncWork(fw *funcWork) {
   111  	fw.ctx = nil
   112  	funcWorkPool.Put(fw)
   113  }
   114  
   115  var funcWorkPool sync.Pool
   116  
   117  type funcWork struct {
   118  	ctx  interface{}
   119  	done chan struct{}
   120  }