github.com/cloudwego/kitex@v0.9.0/pkg/remote/trans/netpoll/trans_server.go

/*
 * Copyright 2021 CloudWeGo Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

// Package netpoll contains server and client implementation for netpoll.

package netpoll

import (
	"context"
	"errors"
	"net"
	"runtime/debug"
	"sync"
	"syscall"

	"github.com/cloudwego/netpoll"

	"github.com/cloudwego/kitex/pkg/klog"
	"github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/remote/trans"
	"github.com/cloudwego/kitex/pkg/utils"
)

// NewTransServerFactory creates a default netpoll transport server factory.
func NewTransServerFactory() remote.TransServerFactory {
	return &netpollTransServerFactory{}
}

type netpollTransServerFactory struct{}

// NewTransServer implements the remote.TransServerFactory interface.
func (f *netpollTransServerFactory) NewTransServer(opt *remote.ServerOption, transHdlr remote.ServerTransHandler) remote.TransServer {
	return &transServer{
		opt:       opt,
		transHdlr: transHdlr,
		lncfg:     trans.NewListenConfig(opt),
	}
}

type transServer struct {
	opt       *remote.ServerOption
	transHdlr remote.ServerTransHandler

	evl       netpoll.EventLoop
	ln        net.Listener
	lncfg     net.ListenConfig
	connCount utils.AtomicInt
	sync.Mutex
}

var _ remote.TransServer = &transServer{}

// CreateListener implements the remote.TransServer interface.
func (ts *transServer) CreateListener(addr net.Addr) (net.Listener, error) {
	if addr.Network() == "unix" {
		syscall.Unlink(addr.String())
	}
	// The network must be "tcp", "tcp4", "tcp6" or "unix".
	ln, err := ts.lncfg.Listen(context.Background(), addr.Network(), addr.String())
	return ln, err
}

// BootstrapServer implements the remote.TransServer interface.
func (ts *transServer) BootstrapServer(ln net.Listener) (err error) {
	if ln == nil {
		return errors.New("listener is nil in netpoll transport server")
	}
	ts.ln = ln
	opts := []netpoll.Option{
		netpoll.WithIdleTimeout(ts.opt.MaxConnectionIdleTime),
		netpoll.WithReadTimeout(ts.opt.ReadWriteTimeout),
	}

	ts.Lock()
	opts = append(opts, netpoll.WithOnPrepare(ts.onConnActive))
	ts.evl, err = netpoll.NewEventLoop(ts.onConnRead, opts...)
	ts.Unlock()

	if err == nil {
		// Convert the listener so that closing it also stops the
		// event loop in netpoll.
		ts.Lock()
		ts.ln, err = netpoll.ConvertListener(ts.ln)
		ts.Unlock()
	}
	if err != nil {
		return err
	}
	return ts.evl.Serve(ts.ln)
}

// Shutdown implements the remote.TransServer interface.
func (ts *transServer) Shutdown() (err error) {
	ts.Lock()
	defer ts.Unlock()

	ctx, cancel := context.WithTimeout(context.Background(), ts.opt.ExitWaitTime)
	defer cancel()
	if g, ok := ts.transHdlr.(remote.GracefulShutdown); ok {
		if ts.ln != nil {
			// 1. stop listener
			ts.ln.Close()

			// 2. signal all active connections to close gracefully
			err = g.GracefulShutdown(ctx)
			if err != nil {
				klog.Warnf("KITEX: server graceful shutdown error: %v", err)
			}
		}
	}
	if ts.evl != nil {
		err = ts.evl.Shutdown(ctx)
	}
	return
}

// ConnCount implements the remote.TransServer interface.
func (ts *transServer) ConnCount() utils.AtomicInt {
	return ts.connCount
}

// TODO: Is it necessary to init RPCInfo when connection is first created?
// Now the strategy is the same as the original one.
// 1. There's no request when the connection is first created.
// 2. Doesn't need to init RPCInfo if it's not RPC request, such as heartbeat.
func (ts *transServer) onConnActive(conn netpoll.Connection) context.Context {
	ctx := context.Background()
	defer transRecover(ctx, conn, "OnActive")
	conn.AddCloseCallback(func(connection netpoll.Connection) error {
		ts.onConnInactive(ctx, conn)
		return nil
	})
	ts.connCount.Inc()
	ctx, err := ts.transHdlr.OnActive(ctx, conn)
	if err != nil {
		ts.onError(ctx, err, conn)
		conn.Close()
		return ctx
	}
	return ctx
}

func (ts *transServer) onConnRead(ctx context.Context, conn netpoll.Connection) error {
	err := ts.transHdlr.OnRead(ctx, conn)
	if err != nil {
		ts.onError(ctx, err, conn)
		if conn != nil {
			// close the connection if OnRead return error
			conn.Close()
		}
	}
	return nil
}

func (ts *transServer) onConnInactive(ctx context.Context, conn netpoll.Connection) {
	defer transRecover(ctx, conn, "OnInactive")
	ts.connCount.Dec()
	ts.transHdlr.OnInactive(ctx, conn)
}

func (ts *transServer) onError(ctx context.Context, err error, conn netpoll.Connection) {
	ts.transHdlr.OnError(ctx, err, conn)
}

func transRecover(ctx context.Context, conn netpoll.Connection, funcName string) {
	panicErr := recover()
	if panicErr != nil {
		if conn != nil {
			klog.CtxErrorf(ctx, "KITEX: panic happened in %s, remoteAddress=%s, error=%v\nstack=%s", funcName, conn.RemoteAddr(), panicErr, string(debug.Stack()))
		} else {
			klog.CtxErrorf(ctx, "KITEX: panic happened in %s, error=%v\nstack=%s", funcName, panicErr, string(debug.Stack()))
		}
	}
}