github.com/15mga/kiwi@v0.0.2-0.20240324021231-b95d5c3ac751/mock/client.go (about)

     1  package mock
     2  
     3  import (
     4  	"fmt"
     5  	"github.com/15mga/kiwi"
     6  	"github.com/15mga/kiwi/graph"
     7  	"github.com/15mga/kiwi/graph/marshall"
     8  	"github.com/15mga/kiwi/network"
     9  	"github.com/15mga/kiwi/util"
    10  	"github.com/15mga/kiwi/worker"
    11  )
    12  
    13  type (
    14  	MsgToStrAny func(client *Client, msg util.IMsg) (point string, data any)
    15  	Decoder     func(kiwi.IAgent, []byte) (svc kiwi.TSvc, mtd kiwi.TCode, msg util.IMsg, err *util.Err)
    16  )
    17  
    18  type Receiver struct {
    19  	Node string
    20  	Fn   MsgToStrAny
    21  }
    22  
    23  type Option struct {
    24  	Id      string
    25  	Name    string
    26  	Ip      string
    27  	Port    int
    28  	Mermaid []byte
    29  	Decoder Decoder
    30  	Head    util.M
    31  	HeadLen uint32
    32  }
    33  
    34  func NewClient(opt Option) (*Client, *util.Err) {
    35  	g := graph.NewGraph("client")
    36  	var ug marshall.Graph
    37  	err := ug.Unmarshall(opt.Mermaid, g)
    38  	if err != nil {
    39  		return nil, err
    40  	}
    41  	m := opt.Head
    42  	if m == nil {
    43  		m = util.M{}
    44  	}
    45  	g.SetData(m)
    46  
    47  	addr := fmt.Sprintf("%s:%d", opt.Ip, opt.Port)
    48  	client := &Client{
    49  		id:            opt.Id,
    50  		decoder:       opt.Decoder,
    51  		graph:         g,
    52  		msgToReceiver: make(map[kiwi.TSvcCode]*Receiver),
    53  	}
    54  	m.Set("client", client)
    55  	dialer := network.NewTcpDialer(opt.Name, addr, client.Receive,
    56  		kiwi.AgentHeadLen(opt.HeadLen),
    57  	)
    58  	dialer.Agent().BindDisconnected(func(agent kiwi.IAgent, err *util.Err) {
    59  		kiwi.Info("disconnect", util.M{
    60  			"addr": agent.Addr(),
    61  		})
    62  	})
    63  	dialer.Agent().BindConnected(func(agent kiwi.IAgent) {
    64  		kiwi.Info("connected", util.M{
    65  			"addr": agent.Addr(),
    66  		})
    67  	})
    68  	e := dialer.Connect(util.Ctx())
    69  	if e != nil {
    70  		return nil, e
    71  	}
    72  	client.dialer = dialer
    73  	return client, nil
    74  }
    75  
    76  type Client struct {
    77  	id            string
    78  	decoder       Decoder
    79  	dialer        kiwi.IDialer
    80  	graph         graph.IGraph
    81  	msgToReceiver map[kiwi.TSvcCode]*Receiver
    82  }
    83  
    84  func (c *Client) Id() string {
    85  	return c.id
    86  }
    87  
    88  func (c *Client) Dialer() kiwi.IDialer {
    89  	return c.dialer
    90  }
    91  
    92  func (c *Client) Graph() graph.IGraph {
    93  	return c.graph
    94  }
    95  
    96  func (c *Client) Receive(agent kiwi.IAgent, bytes []byte) {
    97  	svc, mtd, pkt, err := c.decoder(agent, bytes)
    98  	if err != nil {
    99  		kiwi.Error(err)
   100  		return
   101  	}
   102  	receiver, ok := c.msgToReceiver[kiwi.MergeSvcCode(svc, mtd)]
   103  	if !ok {
   104  		kiwi.Error2(util.EcNotExist, util.M{
   105  			"svc":    svc,
   106  			"method": mtd,
   107  		})
   108  		return
   109  	}
   110  	point, data := receiver.Fn(c, pkt)
   111  	if point == "" {
   112  		return
   113  	}
   114  	if data == nil {
   115  		data = struct{}{}
   116  	}
   117  	nd, err := c.graph.GetNode(receiver.Node)
   118  	if err != nil {
   119  		kiwi.Error(err)
   120  		return
   121  	}
   122  	worker.Share().Push(c.id, c.jobOut, nd, point, data)
   123  }
   124  
   125  func (c *Client) jobOut(params []any) {
   126  	nd, point, data := util.SplitSlc3[graph.INode, string, any](params)
   127  	kiwi.Error(nd.Out(point, data))
   128  }
   129  
   130  func (c *Client) Push(fn util.FnAnySlc, params ...any) {
   131  	worker.Share().Push(c.id, fn, params...)
   132  }
   133  
   134  func (c *Client) BindPointMsg(node, inPoint string, fn graph.MsgToErr) {
   135  	nd, err := c.graph.GetNode(node)
   136  	if err != nil {
   137  		kiwi.Error(err)
   138  		return
   139  	}
   140  	nd.BindFn(inPoint, func(msg graph.IMsg) *util.Err {
   141  		kiwi.Info("process:", msg.ToM())
   142  		return fn(msg)
   143  	})
   144  }
   145  
   146  func (c *Client) BindNetMsg(node string, msg util.IMsg, receiver MsgToStrAny) {
   147  	svc, mtd := kiwi.Codec().MsgToSvcCode(msg)
   148  	c.msgToReceiver[kiwi.MergeSvcCode(svc, mtd)] = &Receiver{
   149  		Node: node,
   150  		Fn:   receiver,
   151  	}
   152  }
   153  
   154  func (c *Client) SetM(key string, data any) {
   155  	c.graph.Data().Set(key, data)
   156  }
   157  
   158  func (c *Client) GetM(key string) (any, bool) {
   159  	return c.graph.Data().Get(key)
   160  }
   161  
   162  func ClientGetM[T any](c *Client, key string) (T, bool) {
   163  	return util.MGet[T](c.graph.Data(), key)
   164  }
   165  
   166  func GraphMsgGetM[T any](msg graph.IMsg, key string) (T, bool) {
   167  	client, _ := util.MGet[*Client](msg.InNode().RootGraph().Data(), "client")
   168  	return ClientGetM[T](client, key)
   169  }
   170  
   171  func MsgGetClient(msg graph.IMsg) *Client {
   172  	client, _ := util.MGet[*Client](msg.InNode().RootGraph().Data(), "client")
   173  	return client
   174  }