gopkg.in/dedis/onet.v2@v2.0.0-20181115163211-c8f3724038a7/simul/monitor/tcpproxy.go (about) 1 // Copyright 2016 The etcd Authors 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 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 package monitor 16 17 import ( 18 "fmt" 19 "io" 20 "math/rand" 21 "net" 22 "sync" 23 "time" 24 25 "gopkg.in/dedis/onet.v2/log" 26 ) 27 28 type remote struct { 29 mu sync.Mutex 30 srv *net.SRV 31 addr string 32 inactive bool 33 } 34 35 func (r *remote) inactivate() { 36 r.mu.Lock() 37 defer r.mu.Unlock() 38 r.inactive = true 39 } 40 41 func (r *remote) tryReactivate() error { 42 conn, err := net.Dial("tcp", r.addr) 43 if err != nil { 44 return err 45 } 46 conn.Close() 47 r.mu.Lock() 48 defer r.mu.Unlock() 49 r.inactive = false 50 return nil 51 } 52 53 func (r *remote) isActive() bool { 54 r.mu.Lock() 55 defer r.mu.Unlock() 56 return !r.inactive 57 } 58 59 // A TCPProxy proxies connections arriving at the Listener to one of 60 // the Endpoints. 61 type TCPProxy struct { 62 Listener net.Listener 63 Endpoints []*net.SRV 64 MonitorInterval time.Duration 65 66 donec chan struct{} 67 68 mu sync.Mutex // guards the following fields 69 remotes []*remote 70 pickCount int // for round robin 71 } 72 73 // Run starts servicing clients. It will not return until Stop is called. 74 func (tp *TCPProxy) Run() error { 75 tp.donec = make(chan struct{}) 76 if tp.MonitorInterval == 0 { 77 tp.MonitorInterval = 5 * time.Minute 78 } 79 for _, srv := range tp.Endpoints { 80 addr := fmt.Sprintf("%s:%d", srv.Target, srv.Port) 81 tp.remotes = append(tp.remotes, &remote{srv: srv, addr: addr}) 82 } 83 84 eps := []string{} 85 for _, ep := range tp.Endpoints { 86 eps = append(eps, fmt.Sprintf("%s:%d", ep.Target, ep.Port)) 87 } 88 89 go tp.runMonitor() 90 for { 91 in, err := tp.Listener.Accept() 92 if err != nil { 93 return err 94 } 95 96 go tp.serve(in) 97 } 98 } 99 100 func (tp *TCPProxy) pick() *remote { 101 var weighted []*remote 102 var unweighted []*remote 103 104 bestPr := uint16(65535) 105 w := 0 106 // find best priority class 107 for _, r := range tp.remotes { 108 switch { 109 case !r.isActive(): 110 case r.srv.Priority < bestPr: 111 bestPr = r.srv.Priority 112 w = 0 113 weighted = nil 114 unweighted = []*remote{r} 115 fallthrough 116 case r.srv.Priority == bestPr: 117 if r.srv.Weight > 0 { 118 weighted = append(weighted, r) 119 w += int(r.srv.Weight) 120 } else { 121 unweighted = append(unweighted, r) 122 } 123 } 124 } 125 if weighted != nil { 126 if len(unweighted) > 0 && rand.Intn(100) == 1 { 127 // In the presence of records containing weights greater 128 // than 0, records with weight 0 should have a very small 129 // chance of being selected. 130 r := unweighted[tp.pickCount%len(unweighted)] 131 tp.pickCount++ 132 return r 133 } 134 // choose a uniform random number between 0 and the sum computed 135 // (inclusive), and select the RR whose running sum value is the 136 // first in the selected order 137 choose := rand.Intn(w) 138 for i := 0; i < len(weighted); i++ { 139 choose -= int(weighted[i].srv.Weight) 140 if choose <= 0 { 141 return weighted[i] 142 } 143 } 144 } 145 if unweighted != nil { 146 for i := 0; i < len(tp.remotes); i++ { 147 picked := tp.remotes[tp.pickCount%len(tp.remotes)] 148 tp.pickCount++ 149 if picked.isActive() { 150 return picked 151 } 152 } 153 } 154 return nil 155 } 156 157 func (tp *TCPProxy) serve(in net.Conn) { 158 var ( 159 err error 160 out net.Conn 161 ) 162 163 for { 164 tp.mu.Lock() 165 remote := tp.pick() 166 tp.mu.Unlock() 167 if remote == nil { 168 break 169 } 170 // TODO: add timeout 171 out, err = net.Dial("tcp", remote.addr) 172 if err == nil { 173 break 174 } 175 remote.inactivate() 176 log.Warnf("deactivated endpoint [%s] due to %v for %v", remote.addr, err, tp.MonitorInterval) 177 } 178 179 if out == nil { 180 in.Close() 181 return 182 } 183 184 go func() { 185 io.Copy(in, out) 186 in.Close() 187 out.Close() 188 }() 189 190 io.Copy(out, in) 191 out.Close() 192 in.Close() 193 } 194 195 func (tp *TCPProxy) runMonitor() { 196 for { 197 select { 198 case <-time.After(tp.MonitorInterval): 199 tp.mu.Lock() 200 for _, rem := range tp.remotes { 201 if rem.isActive() { 202 continue 203 } 204 go func(r *remote) { 205 if err := r.tryReactivate(); err != nil { 206 log.Warnf("failed to activate endpoint [%s] due to %v (stay inactive for another %v)", r.addr, err, tp.MonitorInterval) 207 } else { 208 log.Infof("activated %s", r.addr) 209 } 210 }(rem) 211 } 212 tp.mu.Unlock() 213 case <-tp.donec: 214 return 215 } 216 } 217 } 218 219 // Stop stops a running TCPProxy. After calling Stop, Run will return. 220 func (tp *TCPProxy) Stop() { 221 // graceful shutdown? 222 // shutdown current connections? 223 tp.Listener.Close() 224 close(tp.donec) 225 }