istio.io/istio@v0.0.0-20240520182934-d79c90f27776/pkg/test/echo/server/forwarder/executor.go (about)

     1  // Copyright Istio Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package forwarder
    16  
    17  import (
    18  	"context"
    19  	"fmt"
    20  
    21  	"github.com/hashicorp/go-multierror"
    22  	"go.uber.org/atomic"
    23  	"golang.org/x/sync/semaphore"
    24  )
    25  
    26  const (
    27  	maxConcurrencyPerForward = 20
    28  )
    29  
    30  type executor struct {
    31  	totalRequests  *atomic.Uint64
    32  	activeRequests *atomic.Uint64
    33  	stopCh         chan struct{}
    34  }
    35  
    36  func newExecutor() *executor {
    37  	e := &executor{
    38  		totalRequests:  atomic.NewUint64(0),
    39  		activeRequests: atomic.NewUint64(0),
    40  		stopCh:         make(chan struct{}),
    41  	}
    42  
    43  	return e
    44  }
    45  
    46  func (e *executor) ActiveRequests() uint64 {
    47  	return e.activeRequests.Load()
    48  }
    49  
    50  // NewGroup creates a new group of tasks that can be managed collectively.
    51  // Parallelism is limited by the global maxConcurrency of the executor.
    52  func (e *executor) NewGroup() *execGroup {
    53  	return &execGroup{
    54  		e:   e,
    55  		sem: semaphore.NewWeighted(int64(maxConcurrencyPerForward)),
    56  	}
    57  }
    58  
    59  type execGroup struct {
    60  	e   *executor
    61  	g   multierror.Group
    62  	sem *semaphore.Weighted
    63  }
    64  
    65  // Go runs the given work function asynchronously.
    66  func (g *execGroup) Go(ctx context.Context, work func() error) {
    67  	g.g.Go(func() error {
    68  		g.e.totalRequests.Inc()
    69  		g.e.activeRequests.Inc()
    70  		defer g.e.activeRequests.Dec()
    71  
    72  		// Acquire the group concurrency semaphore.
    73  		if err := g.sem.Acquire(ctx, 1); err != nil {
    74  			return fmt.Errorf("request set timed out: %v", err)
    75  		}
    76  		defer g.sem.Release(1)
    77  
    78  		return work()
    79  	})
    80  }
    81  
    82  func (g *execGroup) Wait() *multierror.Error {
    83  	return g.g.Wait()
    84  }
    85  
    86  func (e *executor) Close() {
    87  	close(e.stopCh)
    88  }