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

     1  package mqtt
     2  
     3  import (
     4  	"time"
     5  
     6  	mqtt "github.com/eclipse/paho.mqtt.golang"
     7  	"github.com/google/uuid"
     8  
     9  	"github.com/machinefi/w3bstream/pkg/depends/base/types"
    10  	conftls "github.com/machinefi/w3bstream/pkg/depends/conf/tls"
    11  	"github.com/machinefi/w3bstream/pkg/depends/x/mapx"
    12  	"github.com/machinefi/w3bstream/pkg/depends/x/misc/retry"
    13  )
    14  
    15  type Broker struct {
    16  	Server        types.Endpoint
    17  	Retry         retry.Retry
    18  	Timeout       types.Duration
    19  	Keepalive     types.Duration
    20  	RetainPublish bool
    21  	QoS           QOS
    22  	Cert          *conftls.X509KeyPair
    23  
    24  	agents *mapx.Map[string, *Client]
    25  }
    26  
    27  func (b *Broker) SetDefault() {
    28  	b.Retry.SetDefault()
    29  	if b.Keepalive == 0 {
    30  		b.Keepalive = types.Duration(3 * time.Hour)
    31  	}
    32  	if b.Timeout == 0 {
    33  		b.Timeout = types.Duration(10 * time.Second)
    34  	}
    35  	if b.Server.IsZero() {
    36  		b.Server.Hostname, b.Server.Port = "127.0.0.1", 1883
    37  	}
    38  	if b.Server.Scheme == "" {
    39  		b.Server.Scheme = "mqtt"
    40  	}
    41  	if b.agents == nil {
    42  		b.agents = mapx.New[string, *Client]()
    43  	}
    44  	if b.QoS > QOS__ONLY_ONCE || b.QoS < 0 {
    45  		b.QoS = QOS__ONCE
    46  	}
    47  }
    48  
    49  func (b *Broker) Init() error {
    50  	if b.Cert != nil {
    51  		if err := b.Cert.Init(); err != nil {
    52  			return err
    53  		}
    54  	}
    55  	return b.Retry.Do(func() error {
    56  		cid := uuid.NewString()
    57  		defer b.CloseByCid(cid)
    58  		_, err := b.Client(cid)
    59  		if err != nil {
    60  			return err
    61  		}
    62  		return nil
    63  	})
    64  }
    65  
    66  func (b *Broker) options(cid string) *mqtt.ClientOptions {
    67  	opt := mqtt.NewClientOptions()
    68  	if cid == "" {
    69  		cid = uuid.NewString()
    70  	}
    71  	opt.SetClientID(cid)
    72  	if !b.Server.IsZero() {
    73  		opt = opt.AddBroker(b.Server.String())
    74  	}
    75  	if b.Server.Username != "" {
    76  		opt.SetUsername(b.Server.Username)
    77  		if b.Server.Password != "" {
    78  			opt.SetPassword(b.Server.Password.String())
    79  		}
    80  	}
    81  	if b.Server.IsTLS() {
    82  		opt.SetTLSConfig(b.Cert.TLSConfig())
    83  	}
    84  
    85  	opt.SetCleanSession(false)
    86  	opt.SetResumeSubs(true)
    87  	opt.SetKeepAlive(b.Keepalive.Duration())
    88  	opt.SetWriteTimeout(b.Timeout.Duration())
    89  	opt.SetConnectTimeout(b.Timeout.Duration())
    90  	return opt
    91  }
    92  
    93  func (b *Broker) Name() string { return "mqtt-broker-cli" }
    94  
    95  func (b *Broker) LivenessCheck() map[string]string {
    96  	m := map[string]string{}
    97  	cid := uuid.NewString()
    98  	defer b.CloseByCid(cid)
    99  	if _, err := b.Client(cid); err != nil {
   100  		m[b.Server.Host()] = err.Error()
   101  		return m
   102  	}
   103  	m[b.Server.Host()] = "ok"
   104  	return m
   105  }
   106  
   107  func (b *Broker) Client(cid string) (*Client, error) {
   108  	return b.ClientWithOptions(b.options(cid))
   109  }
   110  
   111  func (b *Broker) ClientWithOptions(opt *mqtt.ClientOptions) (*Client, error) {
   112  	return b.agents.LoadOrStore(
   113  		opt.ClientID,
   114  		func() (*Client, error) {
   115  			c := &Client{
   116  				cid:    opt.ClientID,
   117  				qos:    b.QoS,
   118  				retain: b.RetainPublish,
   119  				cli:    mqtt.NewClient(opt),
   120  			}
   121  			if err := c.connect(); err != nil {
   122  				return nil, err
   123  			}
   124  			return c, nil
   125  		},
   126  	)
   127  }
   128  
   129  func (b *Broker) Close(c *Client) {
   130  	b.CloseByCid(c.cid)
   131  }
   132  
   133  func (b *Broker) CloseByCid(cid string) {
   134  	if c, ok := b.agents.LoadAndRemove(cid); ok && c != nil {
   135  		c.cli.Disconnect(500)
   136  	}
   137  }