diff --git a/gorush/notification_test.go b/gorush/notification_test.go index 68f6481..3d87457 100644 --- a/gorush/notification_test.go +++ b/gorush/notification_test.go @@ -25,8 +25,6 @@ func TestCorrectConf(t *testing.T) { func TestSenMultipleNotifications(t *testing.T) { PushConf, _ = config.LoadConf("") - InitWorkers(int64(2), 2) - PushConf.Ios.Enabled = true PushConf.Ios.KeyPath = "../certificate/certificate-valid.pem" err := InitAPNSClient() diff --git a/gorush/worker.go b/gorush/worker.go index 872e668..a99c657 100644 --- a/gorush/worker.go +++ b/gorush/worker.go @@ -57,7 +57,9 @@ func queueNotification(req RequestPush) (int, []LogPushEntry) { notification.log = &log notification.AddWaitCount() } - QueueNotification <- *notification + if !tryEnqueue(*notification, QueueNotification) { + LogError.Error("max capacity reached") + } count += len(notification.Tokens) // Count topic message if notification.To != "" { @@ -73,3 +75,15 @@ func queueNotification(req RequestPush) (int, []LogPushEntry) { return count, log } + +// tryEnqueue tries to enqueue a job to the given job channel. Returns true if +// the operation was successful, and false if enqueuing would not have been +// possible without blocking. Job is not enqueued in the latter case. +func tryEnqueue(job PushNotification, jobChan chan<- PushNotification) bool { + select { + case jobChan <- job: + return true + default: + return false + } +} diff --git a/gorush/worker_test.go b/gorush/worker_test.go new file mode 100644 index 0000000..1f22319 --- /dev/null +++ b/gorush/worker_test.go @@ -0,0 +1,17 @@ +package gorush + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestTryEnqueue(t *testing.T) { + chn := make(chan PushNotification, 2) + assert.True(t, tryEnqueue(PushNotification{}, chn)) + assert.Equal(t, 1, len(chn)) + assert.True(t, tryEnqueue(PushNotification{}, chn)) + assert.Equal(t, 2, len(chn)) + assert.False(t, tryEnqueue(PushNotification{}, chn)) + assert.Equal(t, 2, len(chn)) +}