@@ -26,6 +26,7 @@ type RequestPool struct {
2626 semaphore * semaphore.Weighted
2727 lock sync.RWMutex
2828 QueueSize int64
29+ existMap map [string ]bool
2930}
3031
3132type Request struct {
@@ -37,6 +38,7 @@ type Request struct {
3738func (rp * RequestPool ) Start () {
3839 rp .queue = make ([]Request , 0 )
3940 rp .semaphore = semaphore .NewWeighted (rp .QueueSize )
41+ rp .existMap = make (map [string ]bool )
4042}
4143
4244// Submit a request into the pool, returns an error when request is already in the pool
@@ -52,15 +54,15 @@ func (rp *RequestPool) Submit(request []byte) error {
5254 }
5355 rp .lock .Lock ()
5456 defer rp .lock .Unlock ()
55- for _ , existingReq := range rp .queue {
56- if existingReq .ClientID == reqInfo .ClientID && existingReq .ID == reqInfo .ID {
57- rp .semaphore .Release (1 )
58- err := fmt .Sprintf ("a request with ID %v and client ID %v already exists in the pool" , reqInfo .ID , reqInfo .ClientID )
59- rp .Log .Errorf (err )
60- return fmt .Errorf (err )
61- }
57+ existStr := fmt .Sprintf ("%v~%v" ,reqInfo .ClientID , reqInfo .ID )
58+ if _ , exist := rp .existMap [existStr ] ; exist {
59+ rp .semaphore .Release (1 )
60+ err := fmt .Sprintf ("a request with ID %v and client ID %v already exists in the pool" , reqInfo .ID , reqInfo .ClientID )
61+ rp .Log .Errorf (err )
62+ return fmt .Errorf (err )
6263 }
6364 rp .queue = append (rp .queue , req )
65+ rp .existMap [existStr ] = true
6466 return nil
6567}
6668
@@ -78,20 +80,20 @@ func (rp *RequestPool) NextRequests(n int) []Request {
7880func (rp * RequestPool ) RemoveRequest (request Request ) error {
7981 rp .lock .Lock ()
8082 defer rp .lock .Unlock ()
81- removed := false
82- for i , existingReq := range rp .queue {
83- if existingReq .ClientID == request .ClientID && existingReq .ID == request .ID {
84- rp .Log .Infof ("Removed request %v from request pool" , request )
85- rp .queue = append (rp .queue [:i ], rp .queue [i + 1 :]... )
86- removed = true
87- rp .semaphore .Release (1 )
88- break
83+ existStr := fmt .Sprintf ("%v~%v" ,request .ClientID , request .ID )
84+ if _ , exist := rp .existMap [existStr ] ; exist {
85+ for i , existingReq := range rp .queue {
86+ if existingReq .ClientID == request .ClientID && existingReq .ID == request .ID {
87+ rp .Log .Infof ("Removed request %v from request pool" , request )
88+ rp .queue = append (rp .queue [:i ], rp .queue [i + 1 :]... )
89+ delete (rp .existMap , existStr )
90+ rp .semaphore .Release (1 )
91+ return nil
92+ }
8993 }
9094 }
91- if ! removed {
92- err := fmt .Sprintf ("Request %v is not in the pool at remove time" , request )
93- rp .Log .Warnf (err )
94- return fmt .Errorf (err )
95- }
96- return nil
95+ err := fmt .Sprintf ("Request %v is not in the pool at remove time" , request )
96+ rp .Log .Warnf (err )
97+ return fmt .Errorf (err )
98+
9799}
0 commit comments