userspace.go 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172
  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. package tcpproxy
  15. import (
  16. "io"
  17. "net"
  18. "sync"
  19. "time"
  20. "github.com/coreos/pkg/capnslog"
  21. )
  22. var (
  23. plog = capnslog.NewPackageLogger("github.com/coreos/etcd/proxy", "tcpproxy")
  24. )
  25. type remote struct {
  26. mu sync.Mutex
  27. addr string
  28. inactive bool
  29. }
  30. func (r *remote) inactivate() {
  31. r.mu.Lock()
  32. defer r.mu.Unlock()
  33. r.inactive = true
  34. }
  35. func (r *remote) tryReactivate() error {
  36. conn, err := net.Dial("tcp", r.addr)
  37. if err != nil {
  38. return err
  39. }
  40. conn.Close()
  41. r.mu.Lock()
  42. defer r.mu.Unlock()
  43. r.inactive = false
  44. return nil
  45. }
  46. func (r *remote) isActive() bool {
  47. r.mu.Lock()
  48. defer r.mu.Unlock()
  49. return !r.inactive
  50. }
  51. type TCPProxy struct {
  52. Listener net.Listener
  53. Endpoints []string
  54. MonitorInterval time.Duration
  55. donec chan struct{}
  56. mu sync.Mutex // guards the following fields
  57. remotes []*remote
  58. nextRemote int
  59. }
  60. func (tp *TCPProxy) Run() error {
  61. tp.donec = make(chan struct{})
  62. if tp.MonitorInterval == 0 {
  63. tp.MonitorInterval = 5 * time.Minute
  64. }
  65. for _, ep := range tp.Endpoints {
  66. tp.remotes = append(tp.remotes, &remote{addr: ep})
  67. }
  68. go tp.runMonitor()
  69. for {
  70. in, err := tp.Listener.Accept()
  71. if err != nil {
  72. return err
  73. }
  74. go tp.serve(in)
  75. }
  76. }
  77. func (tp *TCPProxy) numRemotes() int {
  78. tp.mu.Lock()
  79. defer tp.mu.Unlock()
  80. return len(tp.remotes)
  81. }
  82. func (tp *TCPProxy) serve(in net.Conn) {
  83. var (
  84. err error
  85. out net.Conn
  86. )
  87. for i := 0; i < tp.numRemotes(); i++ {
  88. remote := tp.pick()
  89. if !remote.isActive() {
  90. continue
  91. }
  92. // TODO: add timeout
  93. out, err = net.Dial("tcp", remote.addr)
  94. if err == nil {
  95. break
  96. }
  97. remote.inactivate()
  98. plog.Warningf("deactivated endpoint [%s] due to %v for %v", remote.addr, err, tp.MonitorInterval)
  99. }
  100. if out == nil {
  101. in.Close()
  102. return
  103. }
  104. go func() {
  105. io.Copy(in, out)
  106. in.Close()
  107. out.Close()
  108. }()
  109. io.Copy(out, in)
  110. out.Close()
  111. in.Close()
  112. }
  113. // pick picks a remote in round-robin fashion
  114. func (tp *TCPProxy) pick() *remote {
  115. tp.mu.Lock()
  116. defer tp.mu.Unlock()
  117. picked := tp.remotes[tp.nextRemote]
  118. tp.nextRemote = (tp.nextRemote + 1) % len(tp.remotes)
  119. return picked
  120. }
  121. func (tp *TCPProxy) runMonitor() {
  122. for {
  123. select {
  124. case <-time.After(tp.MonitorInterval):
  125. tp.mu.Lock()
  126. for _, r := range tp.remotes {
  127. if !r.isActive() {
  128. go func() {
  129. if err := r.tryReactivate(); err != nil {
  130. plog.Warningf("failed to activate endpoint [%s] due to %v (stay inactive for another %v)", r.addr, err, tp.MonitorInterval)
  131. } else {
  132. plog.Printf("activated %s", r.addr)
  133. }
  134. }()
  135. }
  136. }
  137. tp.mu.Unlock()
  138. case <-tp.donec:
  139. return
  140. }
  141. }
  142. }
  143. func (tp *TCPProxy) Stop() {
  144. // graceful shutdown?
  145. // shutdown current connections?
  146. tp.Listener.Close()
  147. close(tp.donec)
  148. }