github.com/lirm/aeron-go@v0.0.0-20230415210743-920325491dc4/aeron/fragmentassembler.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  const (
    27  	DefaultFragmentAssemblyBufferLength = int32(4096)
    28  
    29  	beginFrag    uint8 = 0x80
    30  	endFrag      uint8 = 0x40
    31  	unfragmented uint8 = 0x80 | 0x40
    32  )
    33  
    34  // FragmentAssembler that sits in a chain-of-responsibility pattern that reassembles fragmented messages
    35  // so that the next handler in the chain only sees whole messages.
    36  //
    37  // Unfragmented messages are delegated without copy. Fragmented messages are copied to a temporary
    38  // buffer for reassembly before delegation.
    39  //
    40  // The Header passed to the delegate on assembling a message will be that of the last fragment.
    41  //
    42  // Session based buffers will be allocated and grown as necessary based on the length of messages to be assembled.
    43  type FragmentAssembler struct {
    44  	delegate              term.FragmentHandler
    45  	initialBufferLength   int32
    46  	builderBySessionIdMap map[int32]*bytes.Buffer
    47  }
    48  
    49  // NewFragmentAssembler constructs an adapter to reassemble message fragments and delegate on whole messages.
    50  func NewFragmentAssembler(delegate term.FragmentHandler, initialBufferLength int32) *FragmentAssembler {
    51  	return &FragmentAssembler{
    52  		delegate:              delegate,
    53  		initialBufferLength:   initialBufferLength,
    54  		builderBySessionIdMap: make(map[int32]*bytes.Buffer),
    55  	}
    56  }
    57  
    58  // Clear removes all existing session buffers.
    59  func (f *FragmentAssembler) Clear() {
    60  	for k := range f.builderBySessionIdMap {
    61  		delete(f.builderBySessionIdMap, k)
    62  	}
    63  }
    64  
    65  // OnFragment reassembles and forwards whole messages to the delegate.
    66  func (f *FragmentAssembler) OnFragment(
    67  	buffer *atomic.Buffer,
    68  	offset int32,
    69  	length int32,
    70  	header *logbuffer.Header) {
    71  	flags := header.Flags()
    72  	if (flags & unfragmented) == unfragmented {
    73  		f.delegate(buffer, offset, length, header)
    74  	} else {
    75  		if (flags & beginFrag) == beginFrag {
    76  			builder, ok := f.builderBySessionIdMap[header.SessionId()]
    77  			if !ok {
    78  				builder = &bytes.Buffer{}
    79  				f.builderBySessionIdMap[header.SessionId()] = builder
    80  			}
    81  			builder.Reset()
    82  			buffer.WriteBytes(builder, offset, length)
    83  		} else {
    84  			builder, ok := f.builderBySessionIdMap[header.SessionId()]
    85  			if ok && builder.Len() != 0 {
    86  				buffer.WriteBytes(builder, offset, length)
    87  				if (flags & endFrag) == endFrag {
    88  					msgLength := builder.Len()
    89  					f.delegate(
    90  						atomic.MakeBuffer(builder.Bytes(), msgLength),
    91  						int32(0),
    92  						int32(msgLength),
    93  						header)
    94  					builder.Reset()
    95  				}
    96  			}
    97  		}
    98  	}
    99  }