github.com/siglens/siglens@v0.0.0-20240328180423-f7ce9ae441ed/pkg/utils/semaphore/semaphore.go (about)

     1  /*
     2  Copyright 2023.
     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 semaphore
    18  
    19  import (
    20  	"context"
    21  	"fmt"
    22  	"time"
    23  
    24  	log "github.com/sirupsen/logrus"
    25  	"golang.org/x/sync/semaphore"
    26  )
    27  
    28  var defaultWait = 30 * time.Second
    29  
    30  type WeightedSemaphore struct {
    31  	MaxSize int64  // max size of the semaphore
    32  	Name    string // name of the semaphore
    33  
    34  	sem  *semaphore.Weighted // internal semaphore
    35  	wait time.Duration
    36  }
    37  
    38  /*
    39  Inits a new WeightedSemaphore with max size
    40  
    41  max wait for getting a lock is retryCount * singleTryWait(30s), then will error out
    42  */
    43  func NewDefaultWeightedSemaphore(maxSize int64, name string) *WeightedSemaphore {
    44  	return NewWeightedSemaphore(maxSize, name, defaultWait)
    45  }
    46  
    47  /*
    48  Inits a new WeightedSemaphore with max size
    49  
    50  max wait for getting a lock is retryCount * wait, then will error out
    51  */
    52  func NewWeightedSemaphore(maxSize int64, name string, wait time.Duration) *WeightedSemaphore {
    53  	return &WeightedSemaphore{
    54  		MaxSize: maxSize,
    55  		Name:    name,
    56  		sem:     semaphore.NewWeighted(maxSize),
    57  		wait:    wait,
    58  	}
    59  }
    60  
    61  /*
    62  Wrapper around semaphore.Acquire using contexts and retries for better visibility
    63  
    64  Tries to acquire size from the semaphore. sid is used as a indetifier for log statments
    65  
    66  Returns error if sempahore was not acquired at the end or if never possible to acquire needed size
    67  */
    68  func (s *WeightedSemaphore) TryAcquireWithBackoff(size int64, retryCount int, jobId interface{}) error {
    69  	if size > s.MaxSize {
    70  		return fmt.Errorf("requested %+v, but semaphore was initialized with %+v", size, s.MaxSize)
    71  	}
    72  
    73  	for i := 0; i < retryCount; i++ {
    74  		err := s.acquireSingleJob(size, jobId)
    75  		if err != nil {
    76  			log.Errorf("WeightedSemaphore.%+s Failed to acquire resources for job %+v semaphore after %+v. Retrying %d more times",
    77  				s.Name, jobId, s.wait, retryCount-i-1)
    78  		} else {
    79  			return nil
    80  		}
    81  	}
    82  	return fmt.Errorf("failed to acquire resources for job %+v. requested %+v max size %+v", jobId, size, s.MaxSize)
    83  }
    84  
    85  func (s *WeightedSemaphore) acquireSingleJob(size int64, jobId interface{}) error {
    86  	ctx, cancel := context.WithTimeout(context.Background(), s.wait)
    87  	defer cancel()
    88  	err := s.sem.Acquire(ctx, size)
    89  	return err
    90  }
    91  
    92  /*
    93  Release n resources from the semaphore
    94  */
    95  func (s *WeightedSemaphore) Release(n int64) {
    96  	s.sem.Release(n)
    97  }
    98  
    99  func (s *WeightedSemaphore) TryAcquire(size int64) error {
   100  	if size > s.MaxSize {
   101  		return fmt.Errorf("requested %+v, but semaphore was initialized with %+v", size, s.MaxSize)
   102  	}
   103  	err := s.sem.Acquire(context.TODO(), size)
   104  	if err != nil {
   105  		log.Errorf("WeightedSemaphore.%+s Failed to acquire resources",
   106  			s.Name)
   107  		return fmt.Errorf("failed to acquire resources. Requested %+v max size %+v", size, s.MaxSize)
   108  	} else {
   109  		return nil
   110  	}
   111  }