wehub-resource-sync
129 lines
3.4 KiB
Go
129 lines
3.4 KiB
Go
/*
|
|
* Licensed to the Apache Software Foundation (ASF) under one
|
|
* or more contributor license agreements. See the NOTICE file
|
|
* distributed with this work for additional information
|
|
* regarding copyright ownership. The ASF licenses this file
|
|
* to you under the Apache License, Version 2.0 (the
|
|
* "License"); you may not use this file except in compliance
|
|
* with the License. You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing,
|
|
* software distributed under the License is distributed on an
|
|
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
* KIND, either express or implied. See the License for the
|
|
* specific language governing permissions and limitations
|
|
* under the License.
|
|
*/
|
|
|
|
package queue
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
|
|
"github.com/segmentfault/pacman/log"
|
|
)
|
|
|
|
type Service[T any] interface {
|
|
// Send enqueues a message to be processed asynchronously.
|
|
Send(ctx context.Context, msg T)
|
|
|
|
// RegisterHandler sets the handler function for processing messages.
|
|
RegisterHandler(handler func(ctx context.Context, msg T) error)
|
|
|
|
// Close gracefully shuts down the queue, waiting for pending messages to be processed.
|
|
Close()
|
|
}
|
|
|
|
// Queue is a generic message queue service that processes messages asynchronously.
|
|
// It is thread-safe and supports graceful shutdown.
|
|
type Queue[T any] struct {
|
|
name string
|
|
queue chan T
|
|
handler func(ctx context.Context, msg T) error
|
|
mu sync.RWMutex
|
|
closed bool
|
|
wg sync.WaitGroup
|
|
}
|
|
|
|
// New creates a new queue with the given name and buffer size.
|
|
func New[T any](name string, bufferSize int) *Queue[T] {
|
|
q := &Queue[T]{
|
|
name: name,
|
|
queue: make(chan T, bufferSize),
|
|
}
|
|
q.startWorker()
|
|
return q
|
|
}
|
|
|
|
// Send enqueues a message to be processed asynchronously.
|
|
// It will block if the queue is full.
|
|
func (q *Queue[T]) Send(ctx context.Context, msg T) {
|
|
q.mu.RLock()
|
|
defer q.mu.RUnlock()
|
|
|
|
if q.closed {
|
|
log.Warnf("[%s] queue is closed, dropping message", q.name)
|
|
return
|
|
}
|
|
|
|
select {
|
|
case q.queue <- msg:
|
|
log.Debugf("[%s] enqueued message: %+v", q.name, msg)
|
|
case <-ctx.Done():
|
|
log.Warnf("[%s] context cancelled while sending message", q.name)
|
|
}
|
|
}
|
|
|
|
// RegisterHandler sets the handler function for processing messages.
|
|
// This is thread-safe and can be called at any time.
|
|
func (q *Queue[T]) RegisterHandler(handler func(ctx context.Context, msg T) error) {
|
|
q.mu.Lock()
|
|
defer q.mu.Unlock()
|
|
q.handler = handler
|
|
}
|
|
|
|
// Close gracefully shuts down the queue, waiting for pending messages to be processed.
|
|
func (q *Queue[T]) Close() {
|
|
q.mu.Lock()
|
|
if q.closed {
|
|
q.mu.Unlock()
|
|
return
|
|
}
|
|
q.closed = true
|
|
q.mu.Unlock()
|
|
|
|
close(q.queue)
|
|
q.wg.Wait()
|
|
log.Infof("[%s] queue closed", q.name)
|
|
}
|
|
|
|
// startWorker starts the background goroutine that processes messages.
|
|
func (q *Queue[T]) startWorker() {
|
|
q.wg.Go(func() {
|
|
for msg := range q.queue {
|
|
q.processMessage(msg)
|
|
}
|
|
})
|
|
}
|
|
|
|
// processMessage handles a single message with proper synchronization.
|
|
func (q *Queue[T]) processMessage(msg T) {
|
|
q.mu.RLock()
|
|
handler := q.handler
|
|
q.mu.RUnlock()
|
|
|
|
if handler == nil {
|
|
log.Warnf("[%s] no handler registered, dropping message: %+v", q.name, msg)
|
|
return
|
|
}
|
|
|
|
// Use background context for async processing
|
|
// TODO: Consider adding timeout or using a derived context
|
|
if err := handler(context.TODO(), msg); err != nil {
|
|
log.Errorf("[%s] handler error: %v", q.name, err)
|
|
}
|
|
}
|