code.gitea.io/gitea@v1.22.3/modules/queue/base_channel.go (about)

     1  // Copyright 2023 The Gitea Authors. All rights reserved.
     2  // SPDX-License-Identifier: MIT
     3  
     4  package queue
     5  
     6  import (
     7  	"context"
     8  	"errors"
     9  	"sync"
    10  	"time"
    11  
    12  	"code.gitea.io/gitea/modules/container"
    13  )
    14  
    15  var errChannelClosed = errors.New("channel is closed")
    16  
    17  type baseChannel struct {
    18  	c   chan []byte
    19  	set container.Set[string]
    20  	mu  sync.Mutex
    21  
    22  	isUnique bool
    23  }
    24  
    25  var _ baseQueue = (*baseChannel)(nil)
    26  
    27  func newBaseChannelGeneric(cfg *BaseConfig, unique bool) (baseQueue, error) {
    28  	q := &baseChannel{c: make(chan []byte, cfg.Length), isUnique: unique}
    29  	if unique {
    30  		q.set = container.Set[string]{}
    31  	}
    32  	return q, nil
    33  }
    34  
    35  func newBaseChannelSimple(cfg *BaseConfig) (baseQueue, error) {
    36  	return newBaseChannelGeneric(cfg, false)
    37  }
    38  
    39  func newBaseChannelUnique(cfg *BaseConfig) (baseQueue, error) {
    40  	return newBaseChannelGeneric(cfg, true)
    41  }
    42  
    43  func (q *baseChannel) PushItem(ctx context.Context, data []byte) error {
    44  	if q.c == nil {
    45  		return errChannelClosed
    46  	}
    47  
    48  	if q.isUnique {
    49  		q.mu.Lock()
    50  		has := q.set.Contains(string(data))
    51  		q.mu.Unlock()
    52  		if has {
    53  			return ErrAlreadyInQueue
    54  		}
    55  	}
    56  
    57  	select {
    58  	case q.c <- data:
    59  		if q.isUnique {
    60  			q.mu.Lock()
    61  			q.set.Add(string(data))
    62  			q.mu.Unlock()
    63  		}
    64  		return nil
    65  	case <-time.After(pushBlockTime):
    66  		return context.DeadlineExceeded
    67  	case <-ctx.Done():
    68  		return ctx.Err()
    69  	}
    70  }
    71  
    72  func (q *baseChannel) PopItem(ctx context.Context) ([]byte, error) {
    73  	select {
    74  	case data, ok := <-q.c:
    75  		if !ok {
    76  			return nil, errChannelClosed
    77  		}
    78  		q.mu.Lock()
    79  		q.set.Remove(string(data))
    80  		q.mu.Unlock()
    81  		return data, nil
    82  	case <-ctx.Done():
    83  		return nil, ctx.Err()
    84  	}
    85  }
    86  
    87  func (q *baseChannel) HasItem(ctx context.Context, data []byte) (bool, error) {
    88  	q.mu.Lock()
    89  	defer q.mu.Unlock()
    90  	if !q.isUnique {
    91  		return false, nil
    92  	}
    93  	return q.set.Contains(string(data)), nil
    94  }
    95  
    96  func (q *baseChannel) Len(ctx context.Context) (int, error) {
    97  	q.mu.Lock()
    98  	defer q.mu.Unlock()
    99  
   100  	if q.c == nil {
   101  		return 0, errChannelClosed
   102  	}
   103  
   104  	return len(q.c), nil
   105  }
   106  
   107  func (q *baseChannel) Close() error {
   108  	q.mu.Lock()
   109  	defer q.mu.Unlock()
   110  
   111  	close(q.c)
   112  	if q.isUnique {
   113  		q.set = container.Set[string]{}
   114  	}
   115  
   116  	return nil
   117  }
   118  
   119  func (q *baseChannel) RemoveAll(ctx context.Context) error {
   120  	q.mu.Lock()
   121  	defer q.mu.Unlock()
   122  
   123  	for q.c != nil && len(q.c) > 0 {
   124  		<-q.c
   125  	}
   126  
   127  	if q.isUnique {
   128  		q.set = container.Set[string]{}
   129  	}
   130  	return nil
   131  }