github.com/cloudreve/Cloudreve/v3@v3.0.0-20240224133659-3edb00a6484c/pkg/aria2/rpc/proc.go (about)

     1  package rpc
     2  
     3  import "sync"
     4  
     5  type ResponseProcFn func(resp clientResponse) error
     6  
     7  type ResponseProcessor struct {
     8  	cbs map[uint64]ResponseProcFn
     9  	mu  *sync.RWMutex
    10  }
    11  
    12  func NewResponseProcessor() *ResponseProcessor {
    13  	return &ResponseProcessor{
    14  		make(map[uint64]ResponseProcFn),
    15  		&sync.RWMutex{},
    16  	}
    17  }
    18  
    19  func (r *ResponseProcessor) Add(id uint64, fn ResponseProcFn) {
    20  	r.mu.Lock()
    21  	r.cbs[id] = fn
    22  	r.mu.Unlock()
    23  }
    24  
    25  func (r *ResponseProcessor) remove(id uint64) {
    26  	r.mu.Lock()
    27  	delete(r.cbs, id)
    28  	r.mu.Unlock()
    29  }
    30  
    31  // Process called by recv routine
    32  func (r *ResponseProcessor) Process(resp clientResponse) error {
    33  	id := *resp.Id
    34  	r.mu.RLock()
    35  	fn, ok := r.cbs[id]
    36  	r.mu.RUnlock()
    37  	if ok && fn != nil {
    38  		defer r.remove(id)
    39  		return fn(resp)
    40  	}
    41  	return nil
    42  }