github.com/sunvim/utils@v0.1.0/netpoll/handler.go (about)

     1  // Copyright (c) 2020 Meng Huang (mhboy@outlook.com)
     2  // This package is licensed under a MIT license that can be found in the LICENSE file.
     3  
     4  package netpoll
     5  
     6  import (
     7  	"errors"
     8  	"net"
     9  	"sync"
    10  )
    11  
    12  // ErrHandlerFunc is the error when the HandlerFunc is nil
    13  var ErrHandlerFunc = errors.New("HandlerFunc must be not nil")
    14  
    15  // ErrUpgradeFunc is the error when the Upgrade func is nil
    16  var ErrUpgradeFunc = errors.New("Upgrade function must be not nil")
    17  
    18  // ErrServeFunc is the error when the Serve func is nil
    19  var ErrServeFunc = errors.New("Serve function must be not nil")
    20  
    21  // Context is returned by Upgrade for serving.
    22  type Context interface{}
    23  
    24  // Handler responds to a single request.
    25  type Handler interface {
    26  	// Upgrade upgrades the net.Conn to a Context.
    27  	Upgrade(net.Conn) (Context, error)
    28  	// Serve should serve a single request with the Context.
    29  	Serve(Context) error
    30  }
    31  
    32  // NewHandler returns a new Handler.
    33  func NewHandler(upgrade func(net.Conn) (Context, error), serve func(Context) error) Handler {
    34  	return &ConnHandler{upgrade: upgrade, serve: serve}
    35  }
    36  
    37  // ConnHandler implements the Handler interface.
    38  type ConnHandler struct {
    39  	upgrade func(net.Conn) (Context, error)
    40  	serve   func(Context) error
    41  }
    42  
    43  // SetUpgrade sets the Upgrade function for upgrading the net.Conn.
    44  func (h *ConnHandler) SetUpgrade(upgrade func(net.Conn) (Context, error)) *ConnHandler {
    45  	h.upgrade = upgrade
    46  	return h
    47  }
    48  
    49  // SetServe sets the Serve function for once serving.
    50  func (h *ConnHandler) SetServe(serve func(Context) error) *ConnHandler {
    51  	h.serve = serve
    52  	return h
    53  }
    54  
    55  // Upgrade implements the Handler Upgrade method.
    56  func (h *ConnHandler) Upgrade(conn net.Conn) (Context, error) {
    57  	if h.upgrade == nil {
    58  		return nil, ErrUpgradeFunc
    59  	}
    60  	return h.upgrade(conn)
    61  }
    62  
    63  // Serve implements the Handler Serve method.
    64  func (h *ConnHandler) Serve(ctx Context) error {
    65  	if h.serve == nil {
    66  		return ErrServeFunc
    67  	}
    68  	return h.serve(ctx)
    69  }
    70  
    71  type BytePool interface {
    72  	Get() []byte
    73  	Put(b []byte)
    74  	NumPooled() int
    75  	Width() int
    76  }
    77  
    78  // DataHandler implements the Handler interface.
    79  type DataHandler struct {
    80  	Pool    BytePool
    81  	upgrade func(net.Conn) (net.Conn, error)
    82  	// HandlerFunc is the data Serve function.
    83  	HandlerFunc func(req []byte) (res []byte)
    84  }
    85  
    86  type context struct {
    87  	reading sync.Mutex
    88  	writing sync.Mutex
    89  	upgrade bool
    90  	conn    net.Conn
    91  	pool    BytePool
    92  	buffer  []byte
    93  }
    94  
    95  // SetUpgrade sets the Upgrade function for upgrading the net.Conn.
    96  func (h *DataHandler) SetUpgrade(upgrade func(net.Conn) (net.Conn, error)) {
    97  	h.upgrade = upgrade
    98  }
    99  
   100  // Upgrade sets the net.Conn to a Context.
   101  func (h *DataHandler) Upgrade(conn net.Conn) (Context, error) {
   102  	if h.HandlerFunc == nil {
   103  		return nil, ErrHandlerFunc
   104  	}
   105  	var upgrade bool
   106  	if h.upgrade != nil {
   107  		c, err := h.upgrade(conn)
   108  		if err != nil {
   109  			return nil, err
   110  		} else if c != nil && c != conn {
   111  			upgrade = true
   112  			conn = c
   113  		}
   114  	}
   115  	var ctx = &context{upgrade: upgrade, conn: conn, pool: h.Pool}
   116  	return ctx, nil
   117  }
   118  
   119  // Serve should serve a single request with the Context ctx.
   120  func (h *DataHandler) Serve(ctx Context) error {
   121  	c := ctx.(*context)
   122  	var conn = c.conn
   123  	var n int
   124  	var err error
   125  	var buf []byte
   126  	var req []byte
   127  	buf = c.pool.Get()
   128  	defer c.pool.Put(buf)
   129  	if c.upgrade {
   130  		c.reading.Lock()
   131  	}
   132  	n, err = conn.Read(buf)
   133  	if c.upgrade {
   134  		c.reading.Unlock()
   135  	}
   136  	if err != nil {
   137  		return err
   138  	}
   139  	req = buf[:n]
   140  	res := h.HandlerFunc(req)
   141  	if c.upgrade {
   142  		c.writing.Lock()
   143  	}
   144  	_, err = conn.Write(res)
   145  	if c.upgrade {
   146  		c.writing.Unlock()
   147  	}
   148  	c.pool.Put(res)
   149  	return err
   150  }