github.com/codingeasygo/util@v0.0.0-20231206062002-1ce2f004b7d9/xio/query.go (about)

     1  package xio
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"net"
     7  	"time"
     8  )
     9  
    10  type QueryConn struct {
    11  	sendData chan []byte
    12  	recvData []byte
    13  	recvWait chan int
    14  	err      error
    15  }
    16  
    17  func NewQueryConn() (conn *QueryConn) {
    18  	conn = &QueryConn{
    19  		sendData: make(chan []byte, 1),
    20  		recvWait: make(chan int, 1),
    21  	}
    22  	return
    23  }
    24  
    25  func (q *QueryConn) Read(p []byte) (n int, err error) {
    26  	data := <-q.sendData
    27  	if len(data) < 1 {
    28  		err = q.err
    29  		return
    30  	}
    31  	n = copy(p, data)
    32  	return
    33  }
    34  
    35  func (q *QueryConn) Write(p []byte) (n int, err error) {
    36  	if len(q.recvData) > 0 {
    37  		err = fmt.Errorf("recv twice")
    38  		return
    39  	}
    40  	q.recvData = make([]byte, len(p))
    41  	n = copy(q.recvData, p)
    42  	select {
    43  	case q.recvWait <- 1:
    44  	default:
    45  	}
    46  	return
    47  }
    48  
    49  func (q *QueryConn) Close() (err error) {
    50  	q.err = fmt.Errorf("closed")
    51  	select {
    52  	case q.sendData <- nil:
    53  	default:
    54  	}
    55  	select {
    56  	case q.recvWait <- 1:
    57  	default:
    58  	}
    59  	return
    60  }
    61  
    62  func (q *QueryConn) clearSend() {
    63  	select {
    64  	case <-q.sendData:
    65  	default:
    66  	}
    67  }
    68  
    69  func (q *QueryConn) Query(ctx context.Context, request []byte) (response []byte, err error) {
    70  	select {
    71  	case q.sendData <- request:
    72  		err = q.err
    73  	case <-ctx.Done():
    74  		err = fmt.Errorf("context canceled")
    75  	}
    76  	if err != nil {
    77  		q.clearSend()
    78  		return
    79  	}
    80  	select {
    81  	case <-q.recvWait:
    82  		err = q.err
    83  		response = q.recvData
    84  		q.recvData = nil
    85  	case <-ctx.Done():
    86  		err = fmt.Errorf("context canceled")
    87  	}
    88  	return
    89  }
    90  
    91  // LocalAddr returns the local network address.
    92  func (q *QueryConn) LocalAddr() net.Addr {
    93  	return q
    94  }
    95  
    96  // RemoteAddr returns the remote network address.
    97  func (q *QueryConn) RemoteAddr() net.Addr {
    98  	return q
    99  }
   100  
   101  // SetDeadline for net.Conn
   102  func (q *QueryConn) SetDeadline(t time.Time) error {
   103  	return nil
   104  }
   105  
   106  // SetReadDeadline for net.Conn
   107  func (q *QueryConn) SetReadDeadline(t time.Time) error {
   108  	return nil
   109  }
   110  
   111  // SetWriteDeadline for net.Conn
   112  func (q *QueryConn) SetWriteDeadline(t time.Time) error {
   113  	return nil
   114  }
   115  
   116  // Network is net.Addr impl
   117  func (q *QueryConn) Network() string {
   118  	return "query"
   119  }
   120  
   121  func (q *QueryConn) String() string {
   122  	return "query"
   123  }