github.com/cloudwego/kitex@v0.9.0/pkg/remote/trans/netpoll/conn_extension.go (about)

     1  /*
     2   * Copyright 2021 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package netpoll
    18  
    19  import (
    20  	"context"
    21  	"errors"
    22  	"net"
    23  	"syscall"
    24  
    25  	"github.com/cloudwego/netpoll"
    26  
    27  	"github.com/cloudwego/kitex/pkg/remote"
    28  	"github.com/cloudwego/kitex/pkg/remote/trans"
    29  	"github.com/cloudwego/kitex/pkg/rpcinfo"
    30  )
    31  
    32  // NewNetpollConnExtension to build new netpollConnExtension which implements trans.Extension
    33  func NewNetpollConnExtension() trans.Extension {
    34  	return &netpollConnExtension{}
    35  }
    36  
    37  type netpollConnExtension struct{}
    38  
    39  // SetReadTimeout implements the trans.Extension interface.
    40  func (e *netpollConnExtension) SetReadTimeout(ctx context.Context, conn net.Conn, cfg rpcinfo.RPCConfig, role remote.RPCRole) {
    41  	npConn := conn.(netpoll.Connection)
    42  	if role == remote.Client {
    43  		npConn.SetReadTimeout(trans.GetReadTimeout(cfg))
    44  	} else {
    45  		npConn.SetReadTimeout(cfg.ReadWriteTimeout())
    46  	}
    47  }
    48  
    49  // NewWriteByteBuffer implements the trans.Extension interface.
    50  func (e *netpollConnExtension) NewWriteByteBuffer(ctx context.Context, conn net.Conn, msg remote.Message) remote.ByteBuffer {
    51  	return NewWriterByteBuffer(conn.(netpoll.Connection).Writer())
    52  }
    53  
    54  // NewReadByteBuffer implements the trans.Extension interface.
    55  func (e *netpollConnExtension) NewReadByteBuffer(ctx context.Context, conn net.Conn, msg remote.Message) remote.ByteBuffer {
    56  	return NewReaderByteBuffer(conn.(netpoll.Connection).Reader())
    57  }
    58  
    59  // ReleaseBuffer implements the trans.Extension interface.
    60  func (e *netpollConnExtension) ReleaseBuffer(buffer remote.ByteBuffer, err error) error {
    61  	if buffer != nil {
    62  		return buffer.Release(err)
    63  	}
    64  	return nil
    65  }
    66  
    67  // IsTimeoutErr implements the trans.Extension interface.
    68  func (e *netpollConnExtension) IsTimeoutErr(err error) bool {
    69  	return err != nil && errors.Is(err, netpoll.ErrReadTimeout)
    70  }
    71  
    72  // IsRemoteClosedErr implements the trans.Extension interface.
    73  func (e *netpollConnExtension) IsRemoteClosedErr(err error) bool {
    74  	if err == nil {
    75  		return false
    76  	}
    77  	return errors.Is(err, netpoll.ErrConnClosed) || errors.Is(err, syscall.EPIPE)
    78  }