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

     1  // Copyright 2022 Steven Stern
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  // http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package systests
    16  
    17  import (
    18  	"github.com/lirm/aeron-go/aeron"
    19  	"github.com/lirm/aeron-go/aeron/atomic"
    20  	"github.com/lirm/aeron-go/aeron/logbuffer"
    21  	"github.com/lirm/aeron-go/systests/driver"
    22  	"github.com/stretchr/testify/suite"
    23  	"testing"
    24  	"time"
    25  )
    26  
    27  const (
    28  	streamId           = 1001
    29  	fragmentCountLimit = 10
    30  	messageLength      = 200
    31  )
    32  
    33  type BufferClaimMessageTestSuite struct {
    34  	suite.Suite
    35  	mediaDriver *driver.MediaDriver
    36  	connection  *aeron.Aeron
    37  	errorSink   chan error
    38  	channel     string
    39  }
    40  
    41  func (suite *BufferClaimMessageTestSuite) SetupSuite() {
    42  	mediaDriver, err := driver.StartMediaDriver()
    43  	suite.Require().NoError(err, "Couldn't start Media Driver")
    44  	suite.mediaDriver = mediaDriver
    45  
    46  	suite.errorSink = make(chan error, 100)
    47  	connect, err := aeron.Connect(aeron.
    48  		NewContext().
    49  		AeronDir(suite.mediaDriver.TempDir).
    50  		ErrorHandler(func(err error) { suite.errorSink <- err }))
    51  	if err != nil {
    52  		// Testify does not run TearDownSuite if SetupSuite fails.  We have to manually stop Media Driver.
    53  		suite.mediaDriver.StopMediaDriver()
    54  		suite.Require().NoError(err, "aeron couldn't connect")
    55  	}
    56  	suite.connection = connect
    57  }
    58  
    59  func (suite *BufferClaimMessageTestSuite) TearDownSuite() {
    60  	suite.connection.Close()
    61  	suite.mediaDriver.StopMediaDriver()
    62  }
    63  
    64  func (suite *BufferClaimMessageTestSuite) TestShouldReceivePublishedMessageWithInterleavedAbort() {
    65  	fragmentCount := 0
    66  	fragmentHandler := func(*atomic.Buffer, int32, int32, *logbuffer.Header) {
    67  		fragmentCount++
    68  	}
    69  
    70  	var bufferClaim logbuffer.Claim
    71  	arr := make([]byte, messageLength)
    72  	srcBuffer := atomic.MakeBuffer(arr)
    73  
    74  	subscription, err := suite.connection.AddSubscription(suite.channel, streamId)
    75  	suite.Require().NoError(err)
    76  	defer subscription.Close()
    77  	publication, err := suite.connection.AddPublication(suite.channel, streamId)
    78  	suite.Require().NoError(err)
    79  	defer publication.Close()
    80  
    81  	suite.publishMessage(srcBuffer, *publication)
    82  
    83  	for publication.TryClaim(messageLength, &bufferClaim) < 0 {
    84  		time.Sleep(50 * time.Millisecond)
    85  	}
    86  
    87  	suite.publishMessage(srcBuffer, *publication)
    88  
    89  	bufferClaim.Abort()
    90  
    91  	expectedNumberOfFragments := 2
    92  	numFragments := 0
    93  	for numFragments < expectedNumberOfFragments {
    94  		fragments := subscription.Poll(fragmentHandler, fragmentCountLimit)
    95  		if fragments == 0 {
    96  			time.Sleep(50 * time.Millisecond)
    97  		}
    98  		numFragments += fragments
    99  	}
   100  
   101  	suite.Assert().EqualValues(expectedNumberOfFragments, fragmentCount)
   102  }
   103  
   104  func (suite *BufferClaimMessageTestSuite) TestShouldTransferReservedValue() {
   105  	var bufferClaim logbuffer.Claim
   106  
   107  	subscription, err := suite.connection.AddSubscription(suite.channel, streamId)
   108  	suite.Require().NoError(err)
   109  	defer subscription.Close()
   110  	publication, err := suite.connection.AddPublication(suite.channel, streamId)
   111  	suite.Require().NoError(err)
   112  	defer publication.Close()
   113  
   114  	for publication.TryClaim(messageLength, &bufferClaim) < 0 {
   115  		time.Sleep(50 * time.Millisecond)
   116  	}
   117  
   118  	reservedValue := time.Now().UnixMilli()
   119  	bufferClaim.SetReservedValue(reservedValue)
   120  	bufferClaim.Commit()
   121  
   122  	fragmentHandler := func(_ *atomic.Buffer, _ int32, length int32, header *logbuffer.Header) {
   123  		suite.Assert().EqualValues(messageLength, length)
   124  		suite.Assert().EqualValues(reservedValue, header.GetReservedValue())
   125  	}
   126  
   127  	for 1 != subscription.Poll(fragmentHandler, fragmentCountLimit) {
   128  		time.Sleep(50 * time.Millisecond)
   129  	}
   130  }
   131  
   132  func (suite *BufferClaimMessageTestSuite) publishMessage(srcBuffer *atomic.Buffer, publication aeron.Publication) {
   133  	for publication.Offer(srcBuffer, 0, messageLength, nil) < 0 {
   134  		time.Sleep(50 * time.Millisecond)
   135  	}
   136  }
   137  
   138  func TestBufferClaimMessage(t *testing.T) {
   139  	suite.Run(t, &BufferClaimMessageTestSuite{channel: "aeron:udp?endpoint=localhost:24325"})
   140  	suite.Run(t, &BufferClaimMessageTestSuite{channel: aeron.IpcChannel})
   141  }