trans_pipeline.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 remote

import (
	"context"
	"net"

	"github.com/cloudwego/kitex/pkg/kerrors"
)

// BoundHandler is used to abstract the bound handler.
type BoundHandler interface{}

// OutboundHandler is used to process write event.
type OutboundHandler interface {
	BoundHandler
	Write(ctx context.Context, conn net.Conn, send Message) (context.Context, error)
}

// InboundHandler is used to process read event.
type InboundHandler interface {
	BoundHandler
	OnActive(ctx context.Context, conn net.Conn) (context.Context, error)
	OnInactive(ctx context.Context, conn net.Conn) context.Context
	OnRead(ctx context.Context, conn net.Conn) (context.Context, error)
	OnMessage(ctx context.Context, args, result Message) (context.Context, error)
}

// DuplexBoundHandler can process both inbound and outbound connections.
type DuplexBoundHandler interface {
	OutboundHandler
	InboundHandler
}

// TransPipeline contains TransHandlers.
type TransPipeline struct {
	netHdlr TransHandler

	inboundHdrls  []InboundHandler
	outboundHdrls []OutboundHandler
}

var (
	_ TransHandler       = &TransPipeline{}
	_ ServerTransHandler = &TransPipeline{}
)

func newTransPipeline() *TransPipeline {
	return &TransPipeline{}
}

// NewTransPipeline is used to create a new TransPipeline.
func NewTransPipeline(netHdlr TransHandler) *TransPipeline {
	transPl := newTransPipeline()
	transPl.netHdlr = netHdlr
	netHdlr.SetPipeline(transPl)
	return transPl
}

// AddInboundHandler adds an InboundHandler to the pipeline.
func (p *TransPipeline) AddInboundHandler(hdlr InboundHandler) *TransPipeline {
	p.inboundHdrls = append(p.inboundHdrls, hdlr)
	return p
}

// AddOutboundHandler adds an OutboundHandler to the pipeline.
func (p *TransPipeline) AddOutboundHandler(hdlr OutboundHandler) *TransPipeline {
	p.outboundHdrls = append(p.outboundHdrls, hdlr)
	return p
}

// Write implements the OutboundHandler interface.
func (p *TransPipeline) Write(ctx context.Context, conn net.Conn, sendMsg Message) error {
	var err error
	for _, h := range p.outboundHdrls {
		ctx, err = h.Write(ctx, conn, sendMsg)
		if err != nil {
			return err
		}
	}
	return p.netHdlr.Write(ctx, conn, sendMsg)
}

// OnActive implements the InboundHandler interface.
func (p *TransPipeline) OnActive(ctx context.Context, conn net.Conn) (context.Context, error) {
	var err error
	for _, h := range p.inboundHdrls {
		ctx, err = h.OnActive(ctx, conn)
		if err != nil {
			return ctx, err
		}
	}
	if netHdlr, ok := p.netHdlr.(ServerTransHandler); ok {
		return netHdlr.OnActive(ctx, conn)
	}
	return ctx, nil
}

// OnInactive implements the InboundHandler interface.
func (p *TransPipeline) OnInactive(ctx context.Context, conn net.Conn) {
	for _, h := range p.inboundHdrls {
		ctx = h.OnInactive(ctx, conn)
	}
	p.netHdlr.OnInactive(ctx, conn)
}

// OnRead implements the InboundHandler interface.
func (p *TransPipeline) OnRead(ctx context.Context, conn net.Conn) error {
	var err error
	for _, h := range p.inboundHdrls {
		ctx, err = h.OnRead(ctx, conn)
		if err != nil {
			return err
		}
	}
	if netHdlr, ok := p.netHdlr.(ServerTransHandler); ok {
		return netHdlr.OnRead(ctx, conn)
	}
	return nil
}

// Read reads from conn.
func (p *TransPipeline) Read(ctx context.Context, conn net.Conn, msg Message) error {
	return p.netHdlr.Read(ctx, conn, msg)
}

// OnError calls
func (p *TransPipeline) OnError(ctx context.Context, err error, conn net.Conn) {
	p.netHdlr.OnError(ctx, err, conn)
}

// OnMessage implements the InboundHandler interface.
func (p *TransPipeline) OnMessage(ctx context.Context, args, result Message) (context.Context, error) {
	var err error
	for _, h := range p.inboundHdrls {
		ctx, err = h.OnMessage(ctx, args, result)
		if err != nil {
			return ctx, err
		}
	}
	if result.MessageType() == Exception {
		return ctx, nil
	}
	return p.netHdlr.OnMessage(ctx, args, result)
}

// SetPipeline does nothing now.
func (p *TransPipeline) SetPipeline(transPipe *TransPipeline) {
	// do nothing
}

// GracefulShutdown implements the GracefulShutdown interface.
func (p *TransPipeline) GracefulShutdown(ctx context.Context) error {
	if g, ok := p.netHdlr.(GracefulShutdown); ok {
		return g.GracefulShutdown(ctx)
	}
	return kerrors.ErrNotSupported
}