github.com/gogf/gf@v1.16.9/os/gproc/gproc_comm_receive.go (about)

     1  // Copyright GoFrame Author(https://goframe.org). All Rights Reserved.
     2  //
     3  // This Source Code Form is subject to the terms of the MIT License.
     4  // If a copy of the MIT was not distributed with this file,
     5  // You can obtain one at https://github.com/gogf/gf.
     6  
     7  package gproc
     8  
     9  import (
    10  	"fmt"
    11  	"github.com/gogf/gf/internal/json"
    12  	"net"
    13  
    14  	"github.com/gogf/gf/container/gqueue"
    15  	"github.com/gogf/gf/container/gtype"
    16  	"github.com/gogf/gf/net/gtcp"
    17  	"github.com/gogf/gf/os/gfile"
    18  	"github.com/gogf/gf/os/glog"
    19  	"github.com/gogf/gf/util/gconv"
    20  )
    21  
    22  var (
    23  	// tcpListened marks whether the receiving listening service started.
    24  	tcpListened = gtype.NewBool()
    25  )
    26  
    27  // Receive blocks and receives message from other process using local TCP listening.
    28  // Note that, it only enables the TCP listening service when this function called.
    29  func Receive(group ...string) *MsgRequest {
    30  	// Use atomic operations to guarantee only one receiver goroutine listening.
    31  	if tcpListened.Cas(false, true) {
    32  		go receiveTcpListening()
    33  	}
    34  	var groupName string
    35  	if len(group) > 0 {
    36  		groupName = group[0]
    37  	} else {
    38  		groupName = defaultGroupNameForProcComm
    39  	}
    40  	queue := commReceiveQueues.GetOrSetFuncLock(groupName, func() interface{} {
    41  		return gqueue.New(maxLengthForProcMsgQueue)
    42  	}).(*gqueue.Queue)
    43  
    44  	// Blocking receiving.
    45  	if v := queue.Pop(); v != nil {
    46  		return v.(*MsgRequest)
    47  	}
    48  	return nil
    49  }
    50  
    51  // receiveTcpListening scans local for available port and starts listening.
    52  func receiveTcpListening() {
    53  	var listen *net.TCPListener
    54  	// Scan the available port for listening.
    55  	for i := defaultTcpPortForProcComm; ; i++ {
    56  		addr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("127.0.0.1:%d", i))
    57  		if err != nil {
    58  			continue
    59  		}
    60  		listen, err = net.ListenTCP("tcp", addr)
    61  		if err != nil {
    62  			continue
    63  		}
    64  		// Save the port to the pid file.
    65  		if err := gfile.PutContents(getCommFilePath(Pid()), gconv.String(i)); err != nil {
    66  			panic(err)
    67  		}
    68  		break
    69  	}
    70  	// Start listening.
    71  	for {
    72  		if conn, err := listen.Accept(); err != nil {
    73  			glog.Error(err)
    74  		} else if conn != nil {
    75  			go receiveTcpHandler(gtcp.NewConnByNetConn(conn))
    76  		}
    77  	}
    78  }
    79  
    80  // receiveTcpHandler is the connection handler for receiving data.
    81  func receiveTcpHandler(conn *gtcp.Conn) {
    82  	var (
    83  		result   []byte
    84  		response MsgResponse
    85  	)
    86  	for {
    87  		response.Code = 0
    88  		response.Message = ""
    89  		response.Data = nil
    90  		buffer, err := conn.RecvPkg()
    91  		if len(buffer) > 0 {
    92  			// Package decoding.
    93  			msg := new(MsgRequest)
    94  			if err := json.UnmarshalUseNumber(buffer, msg); err != nil {
    95  				//glog.Error(err)
    96  				continue
    97  			}
    98  			if msg.RecvPid != Pid() {
    99  				// Not mine package.
   100  				response.Message = fmt.Sprintf("receiver pid not match, target: %d, current: %d", msg.RecvPid, Pid())
   101  			} else if v := commReceiveQueues.Get(msg.Group); v == nil {
   102  				// Group check.
   103  				response.Message = fmt.Sprintf("group [%s] does not exist", msg.Group)
   104  			} else {
   105  				// Push to buffer queue.
   106  				response.Code = 1
   107  				v.(*gqueue.Queue).Push(msg)
   108  			}
   109  		} else {
   110  			// Empty package.
   111  			response.Message = "empty package"
   112  		}
   113  		if err == nil {
   114  			result, err = json.Marshal(response)
   115  			if err != nil {
   116  				glog.Error(err)
   117  			}
   118  			if err := conn.SendPkg(result); err != nil {
   119  				glog.Error(err)
   120  			}
   121  		} else {
   122  			// Just close the connection if any error occurs.
   123  			if err := conn.Close(); err != nil {
   124  				glog.Error(err)
   125  			}
   126  			break
   127  		}
   128  	}
   129  }