add sriov token server chain element

pkg/networkservice/common/token/client.go

@@ -1,5 +1,7 @@
 // Copyright (c) 2020 Doc.ai and/or its affiliates.
 //
+// Copyright (c) 2021 Nordix Foundation.
+//
 // SPDX-License-Identifier: Apache-2.0
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
@@ -14,14 +16,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// Package token provides chain element for inserting SRIOV tokens into request
+// Package token provides chain elements for inserting SRIOV tokens into request and response
 package token
 
 import (
 	"context"
 	"os"
 	"strings"
-	"sync"
 
 	"github.com/golang/protobuf/ptypes/empty"
 	"github.com/pkg/errors"
@@ -40,37 +41,21 @@ const (
 )
 
 type tokenClient struct {
-	lock                sync.Mutex
-	tokens              map[string][]string // tokens[tokenName] -> []tokenIDs
-	connectionsByTokens map[string]string   // connectionsByTokens[tokenID] -> connectionID
-	tokensByConnections map[string]string   // tokensByConnections[connectionID] -> tokenID
+	config tokenConfig
 }
 
 // NewClient returns a new token client chain element
 func NewClient() networkservice.NetworkServiceClient {
 	return &tokenClient{
-		tokens:              tokens.FromEnv(os.Environ()),
-		connectionsByTokens: map[string]string{},
-		tokensByConnections: map[string]string{},
+		createTokenElement(tokens.FromEnv(os.Environ())),
 	}
 }
 
 func (c *tokenClient) Request(ctx context.Context, request *networkservice.NetworkServiceRequest, opts ...grpc.CallOption) (*networkservice.Connection, error) {
-	c.lock.Lock()
-	defer c.lock.Unlock()
-
+	var tokenID string
 	if labels := request.GetConnection().GetLabels(); labels != nil {
 		if tokenName, ok := labels[sriovTokenLabel]; ok {
-			var tokenID string
-			for _, tokenID = range c.tokens[tokenName] {
-				if _, ok := c.connectionsByTokens[tokenID]; !ok {
-					c.connectionsByTokens[tokenID] = request.GetConnection().GetId()
-					c.tokensByConnections[request.GetConnection().GetId()] = tokenID
-					break
-				} else {
-					tokenID = ""
-				}
-			}
+			tokenID = c.config.assign(tokenName, request.GetConnection())
 			if tokenID == "" {
 				return nil, errors.Errorf("no free token for the name: %v", tokenName)
 			}
@@ -87,16 +72,17 @@ func (c *tokenClient) Request(ctx context.Context, request *networkservice.Netwo
 			}
 		}
 	}
-	return next.Client(ctx).Request(ctx, request, opts...)
+
+	isEstablished := request.GetConnection().GetNextPathSegment() != nil
+	conn, err := next.Client(ctx).Request(ctx, request, opts...)
+	if err != nil && tokenID != "" && !isEstablished {
+		c.config.release(request.GetConnection())
+	}
+
+	return conn, err
 }
 
 func (c *tokenClient) Close(ctx context.Context, conn *networkservice.Connection, opts ...grpc.CallOption) (*empty.Empty, error) {
-	c.lock.Lock()
-	defer c.lock.Unlock()
-
-	if tokenID, ok := c.tokensByConnections[conn.GetId()]; ok {
-		delete(c.connectionsByTokens, tokenID)
-		delete(c.tokensByConnections, conn.GetId())
-	}
+	c.config.release(conn)
 	return next.Client(ctx).Close(ctx, conn, opts...)
 }

pkg/networkservice/common/token/common.go

@@ -0,0 +1,72 @@
+// Copyright (c) 2021 Nordix Foundation.
+//
+// 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 token
+
+import (
+	"sync"
+
+	"github.com/networkservicemesh/api/pkg/api/networkservice"
+)
+
+type tokenElement struct {
+	lock                sync.Mutex
+	tokens              map[string][]string // tokens[tokenName] -> []tokenIDs
+	connectionsByTokens map[string]string   // connectionsByTokens[tokenID] -> connectionID
+	tokensByConnections map[string]string   // tokensByConnections[connectionID] -> tokenID
+}
+
+type tokenConfig interface {
+	assign(tokenName string, conn *networkservice.Connection) (tokenID string)
+	release(conn *networkservice.Connection)
+}
+
+func createTokenElement(allocatableTokens map[string][]string) tokenConfig {
+	return &tokenElement{tokens: allocatableTokens,
+		connectionsByTokens: map[string]string{},
+		tokensByConnections: map[string]string{}}
+}
+
+func (c *tokenElement) assign(tokenName string, conn *networkservice.Connection) (tokenID string) {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	var ok bool
+	if tokenID, ok = c.tokensByConnections[conn.GetId()]; ok {
+		return tokenID
+	}
+
+	for _, tokenID = range c.tokens[tokenName] {
+		if _, ok := c.connectionsByTokens[tokenID]; !ok {
+			c.connectionsByTokens[tokenID] = conn.GetId()
+			c.tokensByConnections[conn.GetId()] = tokenID
+			break
+		} else {
+			tokenID = ""
+		}
+	}
+	return
+}
+
+func (c *tokenElement) release(conn *networkservice.Connection) {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	if tokenID, ok := c.tokensByConnections[conn.GetId()]; ok {
+		delete(c.connectionsByTokens, tokenID)
+		delete(c.tokensByConnections, conn.GetId())
+	}
+}

pkg/networkservice/common/token/server.go

@@ -0,0 +1,70 @@
+// Copyright (c) 2021 Nordix Foundation.
+//
+// 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 token
+
+import (
+	"context"
+	"os"
+
+	"github.com/golang/protobuf/ptypes/empty"
+	"github.com/networkservicemesh/api/pkg/api/networkservice"
+	"github.com/networkservicemesh/api/pkg/api/networkservice/mechanisms/kernel"
+	"github.com/networkservicemesh/sdk/pkg/networkservice/core/next"
+
+	"github.com/networkservicemesh/sdk-sriov/pkg/networkservice/common/resourcepool"
+	"github.com/networkservicemesh/sdk-sriov/pkg/tools/tokens"
+)
+
+type tokenServer struct {
+	tokenName string
+	config    tokenConfig
+}
+
+// NewServer returns a new token server chain element for the given tokenKey
+func NewServer(tokenKey string) networkservice.NetworkServiceServer {
+	return &tokenServer{
+		tokenName: tokenKey,
+		config: createTokenElement(map[string][]string{
+			tokenKey: tokens.FromEnv(os.Environ())[tokenKey],
+		}),
+	}
+}
+
+func (s *tokenServer) Request(ctx context.Context, request *networkservice.NetworkServiceRequest) (*networkservice.Connection, error) {
+	var tokenID string
+	if mechanism := kernel.ToMechanism(request.GetConnection().GetMechanism()); mechanism != nil || mechanism.Parameters[resourcepool.TokenIDKey] == "" {
+		tokenID = s.config.assign(s.tokenName, request.GetConnection())
+		if tokenID != "" {
+			mechanism.Parameters[resourcepool.TokenIDKey] = tokenID
+		}
+	}
+
+	isEstablished := request.GetConnection().GetNextPathSegment() != nil
+	conn, err := next.Server(ctx).Request(ctx, request)
+	if err != nil && tokenID != "" && !isEstablished {
+		s.config.release(request.GetConnection())
+	}
+
+	return conn, err
+}
+
+func (s *tokenServer) Close(ctx context.Context, conn *networkservice.Connection) (*empty.Empty, error) {
+	if mechanism := kernel.ToMechanism(conn.GetMechanism()); mechanism != nil {
+		s.config.release(conn)
+	}
+	return next.Server(ctx).Close(ctx, conn)
+}

pkg/tools/tokens/tokens.go

@@ -1,5 +1,7 @@
 // Copyright (c) 2020 Doc.ai and/or its affiliates.
 //
+// Copyright (c) 2021 Nordix Foundation.
+//
 // SPDX-License-Identifier: Apache-2.0
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
@@ -23,22 +25,23 @@ import (
 )
 
 const (
-	envPrefix = "NSM_SRIOV_TOKENS_"
+	// EnvPrefix sriov token env name prefix
+	EnvPrefix = "NSM_SRIOV_TOKENS_"
 )
 
 // ToEnv returns a (name, value) pair to store given tokens into the environment variable
 func ToEnv(tokenName string, tokenIDs []string) (name, value string) {
-	return fmt.Sprintf("%s%s", envPrefix, tokenName), strings.Join(tokenIDs, ",")
+	return fmt.Sprintf("%s%s", EnvPrefix, tokenName), strings.Join(tokenIDs, ",")
 }
 
 // FromEnv returns all stored tokens from the list of environment variables
 func FromEnv(envs []string) map[string][]string {
 	tokens := map[string][]string{}
 	for _, env := range envs {
-		if !strings.HasPrefix(env, envPrefix) {
+		if !strings.HasPrefix(env, EnvPrefix) {
 			continue
 		}
-		nameIDs := strings.Split(strings.TrimPrefix(env, envPrefix), "=")
+		nameIDs := strings.Split(strings.TrimPrefix(env, EnvPrefix), "=")
 		tokens[nameIDs[0]] = strings.Split(nameIDs[1], ",")
 	}
 	return tokens