github.com/godaddy-x/freego@v1.0.156/rpcx/pool/conn.go (about) 1 // Copyright 2019 shimingyah. All rights reserved. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // ee the License for the specific language governing permissions and 13 // limitations under the License. 14 15 package pool 16 17 import ( 18 "context" 19 "google.golang.org/grpc" 20 "time" 21 ) 22 23 // Conn single grpc connection inerface 24 type Conn interface { 25 // Value return the actual grpc connection type *grpc.ClientConn. 26 Value() *grpc.ClientConn 27 28 // Close decrease the reference of grpc connection, instead of close it. 29 // if the pool is full, just close it. 30 Close() error 31 32 Context() context.Context 33 34 NewContext(timeout ...time.Duration) 35 } 36 37 // Conn is wrapped grpc.ClientConn. to provide close and value method. 38 type conn struct { 39 cc *grpc.ClientConn 40 pool *pool 41 once bool 42 context context.Context 43 contextCancel context.CancelFunc 44 } 45 46 // Value see Conn interface. 47 func (c *conn) Value() *grpc.ClientConn { 48 return c.cc 49 } 50 51 // Close see Conn interface. 52 func (c *conn) Close() error { 53 c.pool.decrRef() 54 if c.contextCancel != nil { 55 c.contextCancel() 56 c.context = nil 57 c.contextCancel = nil 58 } 59 if c.once { 60 return c.reset() 61 } 62 return nil 63 } 64 65 // Context see Conn interface. 66 func (c *conn) Context() context.Context { 67 return c.context 68 } 69 70 // NewContext see Conn interface. 71 func (c *conn) NewContext(timeout ...time.Duration) { 72 if len(timeout) == 0 { 73 c.context = context.Background() 74 c.contextCancel = nil 75 return 76 } 77 ctx, cancel := context.WithTimeout(context.Background(), timeout[0]) 78 c.context = ctx 79 c.contextCancel = cancel 80 } 81 82 func (c *conn) reset() error { 83 cc := c.cc 84 c.cc = nil 85 c.once = false 86 if cc != nil { 87 return cc.Close() 88 } 89 return nil 90 } 91 92 func (p *pool) wrapConn(cc *grpc.ClientConn, once bool) *conn { 93 return &conn{ 94 cc: cc, 95 pool: p, 96 once: once, 97 } 98 }