github.com/TeaOSLab/EdgeNode@v1.3.8/internal/nodes/origin_conn.go (about)

     1  // Copyright 2023 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
     2  
     3  package nodes
     4  
     5  import (
     6  	teaconst "github.com/TeaOSLab/EdgeNode/internal/const"
     7  	"github.com/TeaOSLab/EdgeNode/internal/goman"
     8  	"github.com/TeaOSLab/EdgeNode/internal/utils/fasttime"
     9  	"github.com/TeaOSLab/EdgeNode/internal/zero"
    10  	"net"
    11  	"sync"
    12  	"time"
    13  )
    14  
    15  const originConnCloseDelaySeconds = 3
    16  
    17  var closingOriginConnMap = map[*OriginConn]zero.Zero{}
    18  var closingOriginConnLocker = &sync.RWMutex{}
    19  
    20  func init() {
    21  	if !teaconst.IsMain {
    22  		return
    23  	}
    24  
    25  	goman.New(func() {
    26  		var ticker = time.NewTicker(originConnCloseDelaySeconds * time.Second)
    27  		for range ticker.C {
    28  			CleanOriginConnsTask()
    29  		}
    30  	})
    31  }
    32  
    33  func CleanOriginConnsTask() {
    34  	var closingConns = []*OriginConn{}
    35  
    36  	closingOriginConnLocker.RLock()
    37  	for conn := range closingOriginConnMap {
    38  		if conn.IsExpired() {
    39  			closingConns = append(closingConns, conn)
    40  		}
    41  	}
    42  	closingOriginConnLocker.RUnlock()
    43  
    44  	if len(closingConns) > 0 {
    45  		for _, conn := range closingConns {
    46  			_ = conn.ForceClose()
    47  			closingOriginConnLocker.Lock()
    48  			delete(closingOriginConnMap, conn)
    49  			closingOriginConnLocker.Unlock()
    50  		}
    51  	}
    52  }
    53  
    54  // OriginConn connection with origin site
    55  type OriginConn struct {
    56  	net.Conn
    57  
    58  	lastReadOk bool
    59  	lastReadAt int64
    60  	isClosed   bool
    61  }
    62  
    63  // NewOriginConn create new origin connection
    64  func NewOriginConn(rawConn net.Conn) net.Conn {
    65  	return &OriginConn{Conn: rawConn}
    66  }
    67  
    68  // Read implement Read() for net.Conn interface
    69  func (this *OriginConn) Read(b []byte) (n int, err error) {
    70  	n, err = this.Conn.Read(b)
    71  	this.lastReadOk = err == nil
    72  	if this.lastReadOk {
    73  		this.lastReadAt = fasttime.Now().Unix()
    74  	}
    75  	return
    76  }
    77  
    78  // Close implement Close() for net.Conn interface
    79  func (this *OriginConn) Close() error {
    80  	if this.lastReadOk && fasttime.Now().Unix()-this.lastReadAt <= originConnCloseDelaySeconds {
    81  		closingOriginConnLocker.Lock()
    82  		closingOriginConnMap[this] = zero.Zero{}
    83  		closingOriginConnLocker.Unlock()
    84  		return nil
    85  	}
    86  
    87  	this.isClosed = true
    88  	return this.Conn.Close()
    89  }
    90  
    91  func (this *OriginConn) ForceClose() error {
    92  	if this.isClosed {
    93  		return nil
    94  	}
    95  
    96  	this.isClosed = true
    97  	return this.Conn.Close()
    98  }
    99  
   100  func (this *OriginConn) IsExpired() bool {
   101  	return fasttime.Now().Unix()-this.lastReadAt > originConnCloseDelaySeconds
   102  }