github.com/cloudwego/kitex@v0.9.0/pkg/generic/descriptor/router.go (about)

     1  /*
     2   * Copyright 2021 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package descriptor
    18  
    19  import (
    20  	"fmt"
    21  	"sync"
    22  )
    23  
    24  // Router http router for bam annotations
    25  type Router interface {
    26  	// Handle register Route to Router
    27  	Handle(rt Route)
    28  	// Lookup FunctionDescriptor from HTTPRequest
    29  	Lookup(req *HTTPRequest) (*FunctionDescriptor, error)
    30  }
    31  
    32  type router struct {
    33  	trees      map[string]*node
    34  	maxParams  uint16
    35  	paramsPool sync.Pool
    36  }
    37  
    38  // NewRouter ...
    39  func NewRouter() Router {
    40  	return &router{}
    41  }
    42  
    43  func (r *router) getParams() *Params {
    44  	ps, _ := r.paramsPool.Get().(*Params)
    45  	ps.params = ps.params[0:0] // reset slice
    46  	return ps
    47  }
    48  
    49  func (r *router) putParams(ps *Params) {
    50  	if ps != nil {
    51  		r.paramsPool.Put(ps)
    52  	}
    53  }
    54  
    55  func (r *router) Handle(rt Route) {
    56  	method := rt.Method()
    57  	path := rt.Path()
    58  	function := rt.Function()
    59  	if method == "" {
    60  		panic("method must not be empty")
    61  	}
    62  	if len(path) < 1 || path[0] != '/' {
    63  		panic("path must begin with '/' in path '" + path + "'")
    64  	}
    65  	if function == nil {
    66  		panic("function descriptor must not be nil")
    67  	}
    68  
    69  	if r.trees == nil {
    70  		r.trees = make(map[string]*node)
    71  	}
    72  
    73  	root := r.trees[method]
    74  	if root == nil {
    75  		root = new(node)
    76  		r.trees[method] = root
    77  	}
    78  
    79  	root.addRoute(path, function)
    80  
    81  	if paramsCount := countParams(path); paramsCount > r.maxParams {
    82  		r.maxParams = paramsCount
    83  	}
    84  
    85  	// Lazy-init paramsPool alloc func
    86  	if r.paramsPool.New == nil && r.maxParams > 0 {
    87  		r.paramsPool.New = func() interface{} {
    88  			ps := Params{
    89  				params:  make([]Param, 0, r.maxParams),
    90  				recycle: r.putParams,
    91  			}
    92  			return &ps
    93  		}
    94  	}
    95  }
    96  
    97  func (r *router) Lookup(req *HTTPRequest) (*FunctionDescriptor, error) {
    98  	root, ok := r.trees[req.GetMethod()]
    99  	if !ok {
   100  		return nil, fmt.Errorf("function lookup failed, no root with method=%s", req.GetMethod())
   101  	}
   102  	fn, ps, _ := root.getValue(req.GetPath(), r.getParams, false)
   103  	if fn == nil {
   104  		r.putParams(ps)
   105  		return nil, fmt.Errorf("function lookup failed, path=%s", req.GetPath())
   106  	}
   107  	req.Params = ps
   108  	return fn, nil
   109  }