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

     1  package mq
     2  
     3  import (
     4  	"context"
     5  	"os"
     6  	"time"
     7  
     8  	"github.com/pkg/errors"
     9  
    10  	"github.com/machinefi/w3bstream/pkg/depends/base/consts"
    11  	"github.com/machinefi/w3bstream/pkg/depends/base/types"
    12  	"github.com/machinefi/w3bstream/pkg/depends/kit/mq"
    13  	"github.com/machinefi/w3bstream/pkg/depends/kit/mq/mem_mq"
    14  	"github.com/machinefi/w3bstream/pkg/depends/x/contextx"
    15  	"github.com/machinefi/w3bstream/pkg/depends/x/misc/must"
    16  )
    17  
    18  type Config struct {
    19  	// Channel worker channel name
    20  	Channel string
    21  	// Store type in memory
    22  	Store StoreType
    23  	// Limit message queue max length
    24  	Limit int
    25  	// WorkerCount task worker count (concurrency)
    26  	WorkerCount      int
    27  	PushQueueTimeout types.Duration
    28  
    29  	*mq.TaskBoard
    30  	*mq.TaskWorker
    31  }
    32  
    33  func (c *Config) SetDefault() {
    34  	if c.Store == STORE_TYPE_UNKNOWN || c.Store > STORE_TYPE__REDIS {
    35  		c.Store = STORE_TYPE__MEM
    36  	}
    37  	if c.PushQueueTimeout == 0 {
    38  		c.PushQueueTimeout = types.Duration(time.Second)
    39  	}
    40  	if c.Limit == 0 {
    41  		c.Limit = 1024
    42  	}
    43  	if c.WorkerCount == 0 {
    44  		c.WorkerCount = 256
    45  	}
    46  }
    47  
    48  func (c *Config) Init() error {
    49  	if c.Channel == "" {
    50  		c.Channel = os.Getenv(consts.EnvProjectName)
    51  	}
    52  
    53  	switch c.Store {
    54  	case STORE_TYPE__MEM:
    55  		tm := mem_mq.New(c.Limit)
    56  		c.TaskWorker = mq.NewTaskWorker(tm, mq.WithWorkerCount(c.WorkerCount), mq.WithChannel(c.Channel))
    57  		c.TaskBoard = mq.NewTaskBoard(tm)
    58  		return nil
    59  	default:
    60  		return errors.New("mq Config init failed")
    61  	}
    62  }
    63  
    64  func (c *Config) Name() string { return "TaskBoard" }
    65  
    66  type kvc struct{}
    67  
    68  func WithMq(ctx context.Context, c *Config) context.Context {
    69  	return contextx.WithValue(ctx, kvc{}, c)
    70  }
    71  
    72  func WithMqContext(c *Config) contextx.WithContext {
    73  	return func(ctx context.Context) context.Context {
    74  		return WithMq(ctx, c)
    75  	}
    76  }
    77  
    78  func MqFromContext(ctx context.Context) (*Config, bool) {
    79  	v, ok := ctx.Value(kvc{}).(*Config)
    80  	return v, ok
    81  }
    82  
    83  func MustMqFromContext(ctx context.Context) *Config {
    84  	v, ok := MqFromContext(ctx)
    85  	must.BeTrue(ok)
    86  	return v
    87  }