Store generated vni in metadata

Signed-off-by: Artem Glazychev <artem.glazychev@xored.com>

pkg/networkservice/common/mechanisms/vxlan/vni/metadata.go

@@ -0,0 +1,51 @@
+// Copyright (c) 2021 Doc.ai and/or its affiliates.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+// 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 vni
+
+import (
+	"context"
+
+	"github.com/networkservicemesh/sdk/pkg/networkservice/utils/metadata"
+)
+
+type key struct{}
+
+func store(ctx context.Context, isClient bool, vni uint32) {
+	metadata.Map(ctx, isClient).Store(key{}, vni)
+}
+
+func delete(ctx context.Context, isClient bool) {
+	metadata.Map(ctx, isClient).Delete(key{})
+}
+
+func load(ctx context.Context, isClient bool) (value uint32, ok bool) {
+	rawValue, ok := metadata.Map(ctx, isClient).Load(key{})
+	if !ok {
+		return
+	}
+	value, ok = rawValue.(uint32)
+	return value, ok
+}
+
+func loadOrStore(ctx context.Context, isClient bool, vni uint32) (value uint32, ok bool) {
+	rawValue, ok := metadata.Map(ctx, isClient).LoadOrStore(key{}, vni)
+	if !ok {
+		return
+	}
+	value, ok = rawValue.(uint32)
+	return value, ok
+}

pkg/networkservice/common/mechanisms/vxlan/vni/server.go

@@ -29,11 +29,14 @@ import (
 	"github.com/networkservicemesh/api/pkg/api/networkservice/mechanisms/vxlan"
 
 	"github.com/networkservicemesh/sdk/pkg/networkservice/core/next"
+	"github.com/networkservicemesh/sdk/pkg/networkservice/utils/metadata"
 )
 
 type vniServer struct {
 	tunnelIP   net.IP
 	tunnelPort uint16
+
+	// This map stores all generated VNIs
 	sync.Map
 }
 
@@ -65,29 +68,38 @@ func (v *vniServer) Request(ctx context.Context, request *networkservice.Network
 	}
 	// If we already have a VNI, make sure we remember it, and go on
 	if k.vni != 0 && mechanism.SrcIP() != nil {
-		_, loaded := v.Map.LoadOrStore(k, &k)
+		_, _ = v.Map.LoadOrStore(k, &k)
+		_, loaded := loadOrStore(ctx, metadata.IsClient(v), k.vni)
 		conn, err := next.Server(ctx).Request(ctx, request)
 		if err != nil && !loaded {
+			delete(ctx, metadata.IsClient(v))
 			v.Map.Delete(k)
 		}
 		return conn, err
 	}
 
-	for {
-		// Generate a random VNI (appropriately odd or even)
-		var err error
-		k.vni, err = mechanism.GenerateRandomVNI()
-		if err != nil {
-			return nil, err
-		}
-		// If its not one already in use, set it and we are good to go
-		if _, ok := v.Map.LoadOrStore(k, &k); !ok {
-			mechanism.SetVNI(k.vni)
-			break
+	if vni, ok := load(ctx, metadata.IsClient(v)); ok {
+		mechanism.SetVNI(vni)
+	} else {
+		for {
+			// Generate a random VNI (appropriately odd or even)
+			var err error
+			k.vni, err = mechanism.GenerateRandomVNI()
+			if err != nil {
+				return nil, err
+			}
+			// If its not one already in use, set it and we are good to go
+			if _, ok := v.Map.LoadOrStore(k, &k); !ok {
+				mechanism.SetVNI(k.vni)
+				store(ctx, metadata.IsClient(v), k.vni)
+				break
+			}
 		}
 	}
+
 	conn, err := next.Server(ctx).Request(ctx, request)
 	if err != nil {
+		delete(ctx, metadata.IsClient(v))
 		v.Map.Delete(k)
 	}
 	return conn, err
@@ -100,6 +112,7 @@ func (v *vniServer) Close(ctx context.Context, conn *networkservice.Connection)
 			vni:         mechanism.VNI(),
 		}
 		if k.vni != 0 && mechanism.SrcIP() != nil {
+			delete(ctx, metadata.IsClient(v))
 			v.Map.LoadAndDelete(k)
 		}
 	}

pkg/networkservice/common/mechanisms/vxlan/vni/server_test.go

@@ -30,6 +30,7 @@ import (
 
 	"github.com/networkservicemesh/sdk/pkg/networkservice/common/mechanisms/vxlan/vni"
 	"github.com/networkservicemesh/sdk/pkg/networkservice/core/next"
+	"github.com/networkservicemesh/sdk/pkg/networkservice/utils/metadata"
 )
 
 func TestVNIServer(t *testing.T) {
@@ -49,7 +50,10 @@ func TestVNIServer(t *testing.T) {
 	}
 	vxlan.ToMechanism(request.GetConnection().GetMechanism()).SetSrcIP(net.ParseIP("192.0.2.1"))
 	var port uint16 = 0
-	server := next.NewNetworkServiceServer(vni.NewServer(net.ParseIP("192.0.2.2")))
+	server := next.NewNetworkServiceServer(
+		metadata.NewServer(),
+		vni.NewServer(net.ParseIP("192.0.2.2")),
+	)
 	conn, err := server.Request(context.Background(), request)
 	assert.Nil(t, err)
 	assert.NotNil(t, conn)
@@ -60,7 +64,10 @@ func TestVNIServer(t *testing.T) {
 	assert.NotEqual(t, 0, mechanism.VNI())
 
 	port = 4466
-	server = next.NewNetworkServiceServer(vni.NewServer(net.ParseIP("192.0.2.2"), vni.WithTunnelPort(port)))
+	server = next.NewNetworkServiceServer(
+		metadata.NewServer(),
+		vni.NewServer(net.ParseIP("192.0.2.2"), vni.WithTunnelPort(port)),
+	)
 	conn, err = server.Request(context.Background(), request)
 	assert.Nil(t, err)
 	assert.NotNil(t, conn)