github.com/ronaksoft/rony@v0.16.26-0.20230807065236-1743dbfe6959/edge/dispatch_context.go (about)

     1  package edge
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"reflect"
     7  	"sync"
     8  
     9  	"github.com/ronaksoft/rony"
    10  	"github.com/ronaksoft/rony/tools"
    11  	"google.golang.org/protobuf/proto"
    12  )
    13  
    14  // DispatchCtx holds the context of the dispatcher's request. Each DispatchCtx could
    15  // hold one or many RequestCtx. DispatchCtx lives until the last of its RequestCtx children.
    16  type DispatchCtx struct {
    17  	edge      *Server
    18  	streamID  int64
    19  	serverID  []byte
    20  	conn      rony.Conn
    21  	req       *rony.MessageEnvelope
    22  	reqFilled bool
    23  	kind      MessageKind
    24  	buf       *tools.LinkedList
    25  	// KeyValue Store Parameters
    26  	mtx sync.RWMutex
    27  	kv  map[string]interface{}
    28  	ctx context.Context
    29  	cf  func()
    30  }
    31  
    32  func newDispatchCtx(edge *Server) *DispatchCtx {
    33  	return &DispatchCtx{
    34  		edge: edge,
    35  		req:  &rony.MessageEnvelope{},
    36  		kv:   make(map[string]interface{}, 3),
    37  		buf:  tools.NewLinkedList(),
    38  	}
    39  }
    40  
    41  func (ctx *DispatchCtx) reset() {
    42  	ctx.reqFilled = false
    43  	for k := range ctx.kv {
    44  		delete(ctx.kv, k)
    45  	}
    46  	ctx.buf.Reset()
    47  }
    48  
    49  func (ctx *DispatchCtx) ServerID() string {
    50  	return string(ctx.serverID)
    51  }
    52  
    53  func (ctx *DispatchCtx) Debug() {
    54  	fmt.Println("###")
    55  	t := reflect.Indirect(reflect.ValueOf(ctx))
    56  	for i := 0; i < t.NumField(); i++ {
    57  		fmt.Println(t.Type().Field(i).Name, t.Type().Field(i).Offset, t.Type().Field(i).Type.Size())
    58  	}
    59  }
    60  
    61  func (ctx *DispatchCtx) Conn() rony.Conn {
    62  	return ctx.conn
    63  }
    64  
    65  func (ctx *DispatchCtx) StreamID() int64 {
    66  	return ctx.streamID
    67  }
    68  
    69  func (ctx *DispatchCtx) FillEnvelope(e *rony.MessageEnvelope) {
    70  	if ctx.reqFilled {
    71  		panic("BUG!!! request has been already filled")
    72  	}
    73  	ctx.reqFilled = true
    74  	e.DeepCopy(ctx.req)
    75  }
    76  
    77  func (ctx *DispatchCtx) Fill(
    78  	requestID uint64, constructor uint64, p proto.Message, kv ...*rony.KeyValue,
    79  ) {
    80  	if ctx.reqFilled {
    81  		panic("BUG!!! request has been already filled")
    82  	}
    83  	ctx.reqFilled = true
    84  	ctx.req.Fill(requestID, constructor, p, kv...)
    85  }
    86  
    87  func (ctx *DispatchCtx) Set(key string, v interface{}) {
    88  	ctx.mtx.Lock()
    89  	ctx.kv[key] = v
    90  	ctx.mtx.Unlock()
    91  }
    92  
    93  func (ctx *DispatchCtx) Get(key string) interface{} {
    94  	ctx.mtx.RLock()
    95  	v := ctx.kv[key]
    96  	ctx.mtx.RUnlock()
    97  
    98  	return v
    99  }
   100  
   101  func (ctx *DispatchCtx) GetBytes(key string, defaultValue []byte) []byte {
   102  	v, ok := ctx.Get(key).([]byte)
   103  	if ok {
   104  		return v
   105  	}
   106  
   107  	return defaultValue
   108  }
   109  
   110  func (ctx *DispatchCtx) GetString(key string, defaultValue string) string {
   111  	v := ctx.Get(key)
   112  	switch x := v.(type) {
   113  	case []byte:
   114  		return tools.ByteToStr(x)
   115  	case string:
   116  		return x
   117  	default:
   118  		return defaultValue
   119  	}
   120  }
   121  
   122  func (ctx *DispatchCtx) GetInt64(key string, defaultValue int64) int64 {
   123  	v, ok := ctx.Get(key).(int64)
   124  	if ok {
   125  		return v
   126  	}
   127  
   128  	return defaultValue
   129  }
   130  
   131  func (ctx *DispatchCtx) GetBool(key string) bool {
   132  	v, ok := ctx.Get(key).(bool)
   133  	if ok {
   134  		return v
   135  	}
   136  
   137  	return false
   138  }
   139  
   140  // Kind identifies that this dispatch context is generated from Tunnel or Gateway. This helps
   141  // developer to apply different strategies based on the source of the incoming message
   142  func (ctx *DispatchCtx) Kind() MessageKind {
   143  	return ctx.kind
   144  }
   145  
   146  func (ctx *DispatchCtx) BufferPush(m *rony.MessageEnvelope) {
   147  	ctx.buf.Append(m)
   148  }
   149  
   150  func (ctx *DispatchCtx) BufferPop(f func(envelope *rony.MessageEnvelope)) bool {
   151  	me, _ := ctx.buf.PickHeadData().(*rony.MessageEnvelope)
   152  	if me == nil {
   153  		return false
   154  	}
   155  	f(me)
   156  	rony.PoolMessageEnvelope.Put(me)
   157  
   158  	return true
   159  }
   160  
   161  func (ctx *DispatchCtx) BufferPopAll(f func(envelope *rony.MessageEnvelope)) {
   162  	for ctx.BufferPop(f) {
   163  	}
   164  }
   165  
   166  func (ctx *DispatchCtx) BufferSize() int32 {
   167  	return ctx.buf.Size()
   168  }
   169  
   170  var dispatchCtxPool = sync.Pool{}
   171  
   172  func acquireDispatchCtx(
   173  	edge *Server, conn rony.Conn,
   174  	streamID int64, serverID []byte, kind MessageKind,
   175  ) *DispatchCtx {
   176  	var ctx *DispatchCtx
   177  	if v := dispatchCtxPool.Get(); v == nil {
   178  		ctx = newDispatchCtx(edge)
   179  	} else {
   180  		ctx = v.(*DispatchCtx)
   181  		ctx.req.Reset()
   182  	}
   183  	ctx.conn = conn
   184  	ctx.kind = kind
   185  	ctx.streamID = streamID
   186  	ctx.serverID = append(ctx.serverID[:0], serverID...)
   187  	ctx.ctx, ctx.cf = context.WithCancel(context.TODO())
   188  
   189  	return ctx
   190  }
   191  
   192  func releaseDispatchCtx(ctx *DispatchCtx) {
   193  	// call cancel func
   194  	ctx.cf()
   195  
   196  	// Reset the Key-Value store
   197  	ctx.reset()
   198  
   199  	// Put back the context into the pool
   200  	dispatchCtxPool.Put(ctx)
   201  }