eintopf.info@v0.13.16/service/oqueue/queue.go (about)

     1  // Copyright (C) 2022 The Eintopf authors
     2  //
     3  // This program is free software: you can redistribute it and/or modify
     4  // it under the terms of the GNU Affero General Public License as
     5  // published by the Free Software Foundation, either version 3 of the
     6  // License, or (at your option) any later version.
     7  //
     8  // This program is distributed in the hope that it will be useful,
     9  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    10  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    11  // GNU Affero General Public License for more details.
    12  //
    13  // You should have received a copy of the GNU Affero General Public License
    14  // along with this program.  If not, see <https://www.gnu.org/licenses/>.
    15  
    16  package oqueue
    17  
    18  import (
    19  	"fmt"
    20  	"log"
    21  	"sync"
    22  	"time"
    23  )
    24  
    25  // Service defines an interface for handling an operator queue. It has methods
    26  // to add and consume operations.
    27  // All methods should be save for concurrent use.
    28  // -go:generate go run github.com/petergtz/pegomock/pegomock generate eintopf.info/service/oqueue Service --output=../../internal/mock/oqueue_service.go --package=mock --mock-name=OQueueService
    29  type Service interface {
    30  	// AddOperation adds an operation to the queue.
    31  	// It might return an error, if the queue is closed.
    32  	AddOperation(op Operation) error
    33  
    34  	// AddSubscriber adds an operation consumer. It gets notified for every
    35  	// operation landing in the queue.
    36  	// If the consumer returns an error, the error gets logged.
    37  	// The amount of worker defines how many operations are to be consumed
    38  	// concurrently.
    39  	AddSubscriber(consume OperationConsumer, worker int)
    40  
    41  	// Close closes the queue. When closed the queue no longer accepts
    42  	// operations.
    43  	// It waits for all operations to be processed before returning.
    44  	Close()
    45  }
    46  
    47  // OperationConsumer is a function, that consumes an operation. The consumer
    48  // might return an error.
    49  type OperationConsumer func(op Operation) error
    50  
    51  // Queue implements an operation queue.
    52  type Queue struct {
    53  	subscriber []*subscriber
    54  	closed     bool
    55  	wgAll      sync.WaitGroup
    56  }
    57  
    58  // Operation is an operation, that can be added to the oqueue.
    59  type Operation interface {
    60  	UserID() string
    61  }
    62  
    63  type subscriber struct {
    64  	queue chan Operation
    65  }
    66  
    67  // NewQueue returns a new queue.
    68  func NewQueue() *Queue {
    69  	return &Queue{
    70  		subscriber: make([]*subscriber, 0),
    71  		closed:     false,
    72  		wgAll:      sync.WaitGroup{},
    73  	}
    74  }
    75  
    76  var ErrClosed = fmt.Errorf("queue was closed")
    77  
    78  // AddOperation adds an operation to the queue.
    79  // Returns ErrClosed if the queue was closed.
    80  func (q *Queue) AddOperation(op Operation) error {
    81  	if q.closed {
    82  		return ErrClosed
    83  	}
    84  	for _, subscriber := range q.subscriber {
    85  		subscriber.queue <- op
    86  	}
    87  	return nil
    88  }
    89  
    90  func (q *Queue) AddSubscriber(consume OperationConsumer, worker int) {
    91  	subscriber := &subscriber{queue: make(chan Operation)}
    92  	q.subscriber = append(q.subscriber, subscriber)
    93  	for i := 0; i < worker; i++ {
    94  		q.wgAll.Add(1)
    95  		go func() {
    96  			for op := range subscriber.queue {
    97  				if err := consume(op); err != nil {
    98  					log.Printf("oqueue: subscriber consume(%v): %s", op, err)
    99  				}
   100  			}
   101  			q.wgAll.Done()
   102  		}()
   103  	}
   104  }
   105  
   106  // Close closes the the queue and waits for all elements in the queue to be
   107  // processed.
   108  func (q *Queue) Close() {
   109  	q.closed = true
   110  
   111  	// Wait for all subscriber queues to be empty before returning.
   112  Outer:
   113  	for {
   114  		time.Sleep(time.Millisecond * 100)
   115  		for _, subscriber := range q.subscriber {
   116  			if len(subscriber.queue) > 1 {
   117  				continue Outer
   118  			}
   119  			close(subscriber.queue)
   120  		}
   121  		break
   122  	}
   123  	q.wgAll.Wait()
   124  }