fastjsonrpc/server.go at main · zc310/fastjsonrpc · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
package fastjsonrpc

import (
	"fmt"

	"github.com/valyala/fasthttp"
	"github.com/valyala/fastjson"
)

func (p *ServerMap) Handler(ctx *fasthttp.RequestCtx) {
	ctx.Response.Header.Set("Content-Type", "application/json; charset=UTF-8")
	defer func() {
		if recover() != nil {
			_, _ = ctx.Write(errInternal)
		}
	}()
	c := getContext()
	c.Ctx = ctx

	p.call(ctx, c)

	_, _ = c.w.WriteTo(ctx)
	putContext(c)
}
func (p *ServerMap) call(ctx *fasthttp.RequestCtx, c *RequestCtx) {
	var err error
	if c.request, err = c.pr.ParseBytes(ctx.PostBody()); err != nil {
		_, _ = c.w.Write(errParse)
		return
	}

	if c.request.Type() == fastjson.TypeArray {
		var a []*fastjson.Value
		a, _ = c.request.Array()
		if len(a) > 32 || len(a) == 0 {
			_, _ = c.w.Write(errInvalidRequest)
			return
		}
		p.batch(a, c)
		return
	}

	if c.request.Type() != fastjson.TypeObject {
		_, _ = c.w.Write(errInvalidRequest)
		return
	}

	c.setRequest(c.request)
	if len(c.Method) == 0 {
		_, _ = c.w.Write(errInvalidRequest)
		return
	}

	f := p.getFun(string(c.Method))
	if f == nil {
		c.Error = errMethodNotFound
		c.writeError(c.w)
		return
	}
	f(c)

	if c.Error == nil {
		c.writeResult(c.w)
	} else {
		c.writeError(c.w)
	}
}
func (p *ServerMap) batch(a []*fastjson.Value, ctx *RequestCtx) {
	bf := getBatchBuffer(len(a))

	for i, sc := range a {
		ct := bf.Ct[i]
		ct.Ctx = ctx.Ctx

		ct.setRequest(sc)
		if ct.request.Type() != fastjson.TypeObject || len(ct.Method) == 0 {
			_, _ = bf.B[i].Write(errInvalidRequest)
			continue
		}
		f := p.getFun(string(ct.Method))
		if f == nil {
			ct.Error = errMethodNotFound
			ct.writeError(bf.B[i])
			continue
		}

		bf.wg.Add(1)

		go func(index int) {
			cc := bf.Ct[index]

			f(cc)

			if cc.Error == nil {
				cc.writeResult(bf.B[index])
			} else {
				cc.writeError(bf.B[index])
			}

			bf.wg.Done()
		}(i)
	}
	bf.wg.Wait()

	_, _ = bf.w.WriteString("[")
	var n int
	for _, b := range bf.B {
		if b.Len() == 0 {
			continue
		}
		if n > 0 {
			_, _ = bf.w.WriteString(",")
		}
		_, _ = b.WriteTo(bf.w)
		n++
	}

	if n > 0 {
		_, _ = fmt.Fprintf(bf.w, "]")
		_, _ = bf.w.WriteTo(ctx.Ctx)
	}

	putBatchBuffer(bf)
}