github.com/lirm/aeron-go@v0.0.0-20230415210743-920325491dc4/aeron/controlledfragmentassembler.go (about)

     1  // Copyright 2016 Stanislav Liberman
     2  // Copyright 2022 Talos, Inc.
     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  package aeron
    17  
    18  import (
    19  	"bytes"
    20  
    21  	"github.com/lirm/aeron-go/aeron/atomic"
    22  	"github.com/lirm/aeron-go/aeron/logbuffer"
    23  	"github.com/lirm/aeron-go/aeron/logbuffer/term"
    24  )
    25  
    26  // ControlledFragmentAssembler that sits in a chain-of-responsibility pattern that reassembles fragmented messages
    27  // so that the next handler in the chain only sees whole messages.
    28  //
    29  // Unfragmented messages are delegated without copy. Fragmented messages are copied to a temporary
    30  // buffer for reassembly before delegation.
    31  //
    32  // The Header passed to the delegate on assembling a message will be that of the last fragment.
    33  //
    34  // Session based buffers will be allocated and grown as necessary based on the length of messages to be assembled.
    35  type ControlledFragmentAssembler struct {
    36  	delegate              term.ControlledFragmentHandler
    37  	initialBufferLength   int32
    38  	builderBySessionIdMap map[int32]*bytes.Buffer
    39  }
    40  
    41  // NewControlledFragmentAssembler constructs an adapter to reassemble message fragments and delegate on whole messages.
    42  func NewControlledFragmentAssembler(delegate term.ControlledFragmentHandler, initialBufferLength int32) *ControlledFragmentAssembler {
    43  	return &ControlledFragmentAssembler{
    44  		delegate:              delegate,
    45  		initialBufferLength:   initialBufferLength,
    46  		builderBySessionIdMap: make(map[int32]*bytes.Buffer),
    47  	}
    48  }
    49  
    50  // OnFragment reassembles and forwards whole messages to the delegate.
    51  func (f *ControlledFragmentAssembler) OnFragment(
    52  	buffer *atomic.Buffer,
    53  	offset int32,
    54  	length int32,
    55  	header *logbuffer.Header) (action term.ControlledPollAction) {
    56  	flags := header.Flags()
    57  	action = term.ControlledPollActionContinue
    58  
    59  	if (flags & unfragmented) == unfragmented {
    60  		action = f.delegate(buffer, offset, length, header)
    61  	} else {
    62  		if (flags & beginFrag) == beginFrag {
    63  			builder, ok := f.builderBySessionIdMap[header.SessionId()]
    64  			if !ok {
    65  				builder = &bytes.Buffer{}
    66  				f.builderBySessionIdMap[header.SessionId()] = builder
    67  			}
    68  			builder.Reset()
    69  			buffer.WriteBytes(builder, offset, length)
    70  		} else {
    71  			builder, ok := f.builderBySessionIdMap[header.SessionId()]
    72  			if ok {
    73  				if limit := builder.Len(); limit > 0 {
    74  					buffer.WriteBytes(builder, offset, length)
    75  					if (flags & endFrag) == endFrag {
    76  						msgLength := builder.Len()
    77  						action := f.delegate(
    78  							atomic.MakeBuffer(builder.Bytes(), msgLength),
    79  							int32(0),
    80  							int32(msgLength),
    81  							header)
    82  						if action == term.ControlledPollActionAbort {
    83  							builder.Truncate(limit)
    84  						} else {
    85  							builder.Reset()
    86  						}
    87  					}
    88  				}
    89  			}
    90  		}
    91  	}
    92  	return
    93  }