github.com/machinefi/w3bstream@v1.6.5-rc9.0.20240426031326-b8c7c4876e72/pkg/depends/kit/mq/mem_mq/task_test.go (about)

     1  package mem_mq_test
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"fmt"
     7  	"sync"
     8  	"testing"
     9  
    10  	"github.com/machinefi/w3bstream/pkg/depends/kit/mq"
    11  	"github.com/machinefi/w3bstream/pkg/depends/kit/mq/mem_mq"
    12  	"github.com/machinefi/w3bstream/pkg/depends/kit/mq/worker"
    13  )
    14  
    15  var (
    16  	tm = mem_mq.New(10000)
    17  	ch = "mm"
    18  )
    19  
    20  func NewTask(subject, id string, payload ...interface{}) *Task {
    21  	t := &Task{}
    22  	t.SetSubject(subject)
    23  
    24  	if id != "" {
    25  		t.SetID(id)
    26  	}
    27  
    28  	if len(payload) == 0 {
    29  		return t
    30  	}
    31  	if pl, ok := payload[0].([]byte); ok {
    32  		t.Write(pl)
    33  	}
    34  
    35  	return t
    36  }
    37  
    38  type Task struct {
    39  	mq.TaskHeader
    40  	bytes.Buffer
    41  }
    42  
    43  func (t *Task) Payload() []byte { return t.Bytes() }
    44  
    45  func BenchmarkTaskManager(b *testing.B) {
    46  	for i := 0; i < b.N; i++ {
    47  		_ = tm.Push(ch, NewTask("", fmt.Sprintf("%d", i), nil))
    48  		_, _ = tm.Pop(ch)
    49  	}
    50  }
    51  
    52  func TestTaskManager(t *testing.T) {
    53  	_ = tm.Clear(ch)
    54  
    55  	for i := 0; i < 1000; i++ {
    56  		_ = tm.Push(ch, NewTask("", fmt.Sprintf("%d", i), nil))
    57  		_ = tm.Push(ch, NewTask("", fmt.Sprintf("%d", i), nil))
    58  		_ = tm.Push(ch, NewTask("", fmt.Sprintf("%d", i), nil))
    59  		_ = tm.Push(ch, NewTask("", fmt.Sprintf("%d", i), nil))
    60  		_ = tm.Push(ch, NewTask("", fmt.Sprintf("%d", i), nil))
    61  	}
    62  
    63  	wg := sync.WaitGroup{}
    64  	wg.Add(1000)
    65  
    66  	w := worker.New(func(ctx context.Context) error {
    67  		task, err := tm.Pop(ch)
    68  		if err != nil {
    69  			return err
    70  		}
    71  		if task == nil {
    72  			return nil
    73  		}
    74  		wg.Add(-1)
    75  		return nil
    76  	}, 10)
    77  
    78  	go w.Start(context.Background())
    79  	wg.Wait()
    80  }