<commit_before>/*=========================================================================

  Program:   Insight Segmentation & Registration Toolkit
  Module:    itkMeshSpatialObjectTest.cxx
  Language:  C++
  Date:      $Date$
  Version:   $Revision$

  Copyright (c) Insight Software Consortium. All rights reserved.
  See ITKCopyright.txt or http://www.itk.org/HTML/Copyright.htm for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notices for more information.

=========================================================================*/
#if defined(_MSC_VER)
#pragma warning ( disable : 4786 )
#endif

#include <itkDefaultDynamicMeshTraits.h>
#include <itkMesh.h>
#include <itkMeshSpatialObject.h>
#include <itkTetrahedronCell.h>

int itkMeshSpatialObjectTest(int, char * [] )
{
  typedef itk::DefaultDynamicMeshTraits< float , 3, 3 > MeshTrait;
  typedef itk::Mesh<float,3,MeshTrait>                  MeshType;
  typedef MeshType::CellTraits                          CellTraits;
  typedef itk::CellInterface< float, CellTraits >       CellInterfaceType;
  typedef itk::TetrahedronCell<CellInterfaceType>       TetraCellType;
  typedef itk::MeshSpatialObject<MeshType>              MeshSpatialObjectType;
  typedef MeshType::PointType                           PointType;
  typedef MeshType::CellType                            CellType;
  typedef CellType::CellAutoPointer                     CellAutoPointer;

  // Create an itkMesh
  MeshType::Pointer mesh = MeshType::New();

  MeshType::CoordRepType testPointCoords[4][3]
    = { {0,0,0}, {9,0,0}, {9,9,0}, {0,0,9} };

  unsigned long tetraPoints[4] = {0,1,2,4};

  int i;
  for(i=0; i < 4 ; ++i)
    {
    mesh->SetPoint(i, PointType(testPointCoords[i]));
    }

  mesh->SetCellsAllocationMethod( MeshType::CellsAllocatedDynamicallyCellByCell );
  CellAutoPointer testCell1;
  testCell1.TakeOwnership(  new TetraCellType );
  testCell1->SetPointIds(tetraPoints);
  mesh->SetCell(0, testCell1 );

  // Create the mesh Spatial Object
  MeshSpatialObjectType::Pointer meshSO = MeshSpatialObjectType::New();
  meshSO->SetMesh(mesh);

  std::cout << "Testing GetMesh(): ";

  if(mesh != meshSO->GetMesh())
    {
    std::cout<<"[FAILED]"<<std::endl;
    return EXIT_FAILURE;
    }
  std::cout<<"[PASSED]"<<std::endl;

  std::cout << "Testing Bounding Box: ";

  if( (meshSO->GetBoundingBox()->GetBounds()[0] != 0)
   || (meshSO->GetBoundingBox()->GetBounds()[1] != 9)
   || (meshSO->GetBoundingBox()->GetBounds()[2] != 0)
   || (meshSO->GetBoundingBox()->GetBounds()[3] != 9)
   || (meshSO->GetBoundingBox()->GetBounds()[4] != 0)
   || (meshSO->GetBoundingBox()->GetBounds()[5] != 9)
   )
    {
    std::cout<<"[FAILED]"<<std::endl;
    return EXIT_FAILURE;
    }

  std::cout<<"[PASSED]"<<std::endl;

  // Testing is inside
  std::cout << "Testing IsInside: ";

  MeshSpatialObjectType::PointType inside;
  inside[0] = 1;
  inside[1] = 1;
  inside[2] = 1;
  MeshSpatialObjectType::PointType outside;
  outside[0] = 0;
  outside[1] = 3;
  outside[2] = 0;

  if(!meshSO->IsInside(inside) || meshSO->IsInside(outside))
    {
    std::cout<<"[FAILED]"<<std::endl;
    return EXIT_FAILURE;
    }
  std::cout<<"[PASSED]"<<std::endl;

  // Testing is valueAt
  std::cout << "Testing ValueAt: ";
  double value;
  meshSO->ValueAt(inside,value);
  if(!value)
    {
    std::cout<<"[FAILED]"<<std::endl;
    return EXIT_FAILURE;
    }

  meshSO->ValueAt(outside,value);
  if(value)
    {
    std::cout<<"[FAILED]"<<std::endl;
    return EXIT_FAILURE;
    }

  std::cout<<"[PASSED]"<<std::endl;

  // Testing IsInside() for triangle mesh
  std::cout << "Testing IsInside() Triangle: ";
  typedef itk::TriangleCell<CellInterfaceType>          TriangleCellType;

  // Create an itkMesh
  MeshType::Pointer meshTriangle = MeshType::New();

  MeshType::CoordRepType testTrianglePointCoords[4][3]
    = { {50,50,64}, {50,100,64}, {100,50,64} , {100,100,64}};

  unsigned long trianglePoint1[3] = {0,1,2};

  for(i=0; i < 4 ; ++i)
    {
    meshTriangle->SetPoint(i, PointType(testTrianglePointCoords[i]));
    }

  unsigned long trianglePoint2[] = {1,2,3};

  meshTriangle->SetCellsAllocationMethod( MeshType::CellsAllocatedDynamicallyCellByCell );
  CellAutoPointer testCell3;
  testCell3.TakeOwnership(  new TriangleCellType );
  testCell3->SetPointIds(trianglePoint1);
  meshTriangle->SetCell(0, testCell3 );

  CellAutoPointer testCell4;
  testCell4.TakeOwnership(  new TriangleCellType );
  testCell4->SetPointIds(trianglePoint2);
  meshTriangle->SetCell(1, testCell4 );

  // Create the mesh Spatial Object
  MeshSpatialObjectType::Pointer meshTriangleSO = MeshSpatialObjectType::New();
  meshTriangleSO->SetMesh(meshTriangle);

  itk::Point<double,3> pIn;
  pIn[0] = 60;
  pIn[1] = 60;
  pIn[2] = 64;
  itk::Point<double,3> pOut;
  pOut[0] = 60;
  pOut[1] = 102;
  pOut[2] = 64;
  if(!meshTriangleSO->IsInside(pIn) || meshTriangleSO->IsInside(pOut))
    {
    std::cout<<"[FAILED]"<<std::endl;
    return EXIT_FAILURE;
    }
  std::cout<<"[PASSED]"<<std::endl;

  std::cout<<"[TEST DONE]"<<std::endl;

  return EXIT_SUCCESS;

}
<commit_msg>BUG: bad index for tetraPoints.<commit_after>/*=========================================================================

  Program:   Insight Segmentation & Registration Toolkit
  Module:    itkMeshSpatialObjectTest.cxx
  Language:  C++
  Date:      $Date$
  Version:   $Revision$

  Copyright (c) Insight Software Consortium. All rights reserved.
  See ITKCopyright.txt or http://www.itk.org/HTML/Copyright.htm for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notices for more information.

=========================================================================*/
#if defined(_MSC_VER)
#pragma warning ( disable : 4786 )
#endif

#include <itkDefaultDynamicMeshTraits.h>
#include <itkMesh.h>
#include <itkMeshSpatialObject.h>
#include <itkTetrahedronCell.h>

int itkMeshSpatialObjectTest(int, char * [] )
{
  typedef itk::DefaultDynamicMeshTraits< float , 3, 3 > MeshTrait;
  typedef itk::Mesh<float,3,MeshTrait>                  MeshType;
  typedef MeshType::CellTraits                          CellTraits;
  typedef itk::CellInterface< float, CellTraits >       CellInterfaceType;
  typedef itk::TetrahedronCell<CellInterfaceType>       TetraCellType;
  typedef itk::MeshSpatialObject<MeshType>              MeshSpatialObjectType;
  typedef MeshType::PointType                           PointType;
  typedef MeshType::CellType                            CellType;
  typedef CellType::CellAutoPointer                     CellAutoPointer;

  // Create an itkMesh
  MeshType::Pointer mesh = MeshType::New();

  MeshType::CoordRepType testPointCoords[4][3]
    = { {0,0,0}, {9,0,0}, {9,9,0}, {0,0,9} };

  unsigned long tetraPoints[4] = {0,1,2,3};

  int i;
  for(i=0; i < 4 ; ++i)
    {
    mesh->SetPoint(i, PointType(testPointCoords[i]));
    }

  mesh->SetCellsAllocationMethod( MeshType::CellsAllocatedDynamicallyCellByCell );
  CellAutoPointer testCell1;
  testCell1.TakeOwnership(  new TetraCellType );
  testCell1->SetPointIds(tetraPoints);
  mesh->SetCell(0, testCell1 );

  // Create the mesh Spatial Object

  MeshSpatialObjectType::Pointer meshSO = MeshSpatialObjectType::New();
  meshSO->SetMesh(mesh);

  std::cout << "Testing GetMesh(): ";

  if(mesh != meshSO->GetMesh())
    {
    std::cout<<"[FAILED]"<<std::endl;
    return EXIT_FAILURE;
    }
  std::cout<<"[PASSED]"<<std::endl;

  std::cout << "Testing Bounding Box: ";

  if( (meshSO->GetBoundingBox()->GetBounds()[0] != 0)
   || (meshSO->GetBoundingBox()->GetBounds()[1] != 9)
   || (meshSO->GetBoundingBox()->GetBounds()[2] != 0)
   || (meshSO->GetBoundingBox()->GetBounds()[3] != 9)
   || (meshSO->GetBoundingBox()->GetBounds()[4] != 0)
   || (meshSO->GetBoundingBox()->GetBounds()[5] != 9)
   )
    {
    std::cout<<"[FAILED]"<<std::endl;
    return EXIT_FAILURE;
    }

  std::cout<<"[PASSED]"<<std::endl;

  // Testing is inside
  std::cout << "Testing IsInside: ";

  MeshSpatialObjectType::PointType inside;
  inside[0] = 1;
  inside[1] = 1;
  inside[2] = 1;
  MeshSpatialObjectType::PointType outside;
  outside[0] = 0;
  outside[1] = 3;
  outside[2] = 0;

  if(!meshSO->IsInside(inside) || meshSO->IsInside(outside))
    {
    std::cout<<"[FAILED]"<<std::endl;
    return EXIT_FAILURE;
    }
  std::cout<<"[PASSED]"<<std::endl;

  // Testing is valueAt
  std::cout << "Testing ValueAt: ";
  double value;
  meshSO->ValueAt(inside,value);
  if(!value)
    {
    std::cout<<"[FAILED]"<<std::endl;
    return EXIT_FAILURE;
    }

  meshSO->ValueAt(outside,value);
  if(value)
    {
    std::cout<<"[FAILED]"<<std::endl;
    return EXIT_FAILURE;
    }

  std::cout<<"[PASSED]"<<std::endl;

  // Testing IsInside() for triangle mesh
  std::cout << "Testing IsInside() Triangle: ";
  typedef itk::TriangleCell<CellInterfaceType>          TriangleCellType;

  // Create an itkMesh
  MeshType::Pointer meshTriangle = MeshType::New();

  MeshType::CoordRepType testTrianglePointCoords[4][3]
    = { {50,50,64}, {50,100,64}, {100,50,64} , {100,100,64}};

  unsigned long trianglePoint1[3] = {0,1,2};

  for(i=0; i < 4 ; ++i)
    {
    meshTriangle->SetPoint(i, PointType(testTrianglePointCoords[i]));
    }

  unsigned long trianglePoint2[] = {1,2,3};

  meshTriangle->SetCellsAllocationMethod( MeshType::CellsAllocatedDynamicallyCellByCell );
  CellAutoPointer testCell3;
  testCell3.TakeOwnership(  new TriangleCellType );
  testCell3->SetPointIds(trianglePoint1);
  meshTriangle->SetCell(0, testCell3 );

  CellAutoPointer testCell4;
  testCell4.TakeOwnership(  new TriangleCellType );
  testCell4->SetPointIds(trianglePoint2);
  meshTriangle->SetCell(1, testCell4 );

  // Create the mesh Spatial Object
  MeshSpatialObjectType::Pointer meshTriangleSO = MeshSpatialObjectType::New();
  meshTriangleSO->SetMesh(meshTriangle);

  itk::Point<double,3> pIn;
  pIn[0] = 60;
  pIn[1] = 60;
  pIn[2] = 64;
  itk::Point<double,3> pOut;
  pOut[0] = 60;
  pOut[1] = 102;
  pOut[2] = 64;
  if(!meshTriangleSO->IsInside(pIn) || meshTriangleSO->IsInside(pOut))
    {
    std::cout<<"[FAILED]"<<std::endl;
    return EXIT_FAILURE;
    }
  std::cout<<"[PASSED]"<<std::endl;

  std::cout<<"[TEST DONE]"<<std::endl;

  return EXIT_SUCCESS;

}
<|endoftext|>


<commit_before>/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

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.
==============================================================================*/

#include "tensorflow/core/profiler/utils/kernel_stats_utils.h"

#include <algorithm>
#include <string>
#include <tuple>
#include <vector>

#include "absl/algorithm/container.h"
#include "absl/strings/match.h"
#include "absl/strings/numbers.h"
#include "absl/strings/str_split.h"
#include "absl/strings/string_view.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/profiler/protobuf/kernel_stats.pb.h"

namespace tensorflow {
namespace profiler {

namespace {

// The maximum number of Kernels displayed on Kernel Stats page.
const int kMaxNumOfKernels = 1000;

}  // namespace

void ParseKernelLaunchParams(absl::string_view xstat_kernel_details,
                             KernelReport* kernel) {
  const std::vector<absl::string_view> params =
      absl::StrSplit(xstat_kernel_details, absl::ByAnyChar(" \n"));

  constexpr uint32 kNumDimensions = 3;
  for (uint32 dim = 0; dim < kNumDimensions; ++dim) {
    kernel->add_block_dim(1);
    kernel->add_grid_dim(1);
  }

  // Process tokens.
  for (const auto& param : params) {
    const std::vector<absl::string_view> key_value = absl::StrSplit(param, ':');
    if (key_value.size() != 2) {
      // Unrecognized token.
      continue;
    }
    absl::string_view key = key_value[0];
    absl::string_view value_str = key_value[1];
    uint32 value = 0;
    double pct = 0.0;
    // Cases that consume a pair of tokens "key:value".
    if (key == "regs" && absl::SimpleAtoi(value_str, &value)) {
      kernel->set_registers_per_thread(value);
    } else if (key == "static_shared" && absl::SimpleAtoi(value_str, &value)) {
      kernel->set_static_shmem_bytes(value);
    } else if (key == "dynamic_shared" && absl::SimpleAtoi(value_str, &value)) {
      kernel->set_dynamic_shmem_bytes(value);
    } else if (key == "block") {
      const std::vector<absl::string_view>& block =
          absl::StrSplit(value_str, ',');
      uint32 tmp[3];
      if (block.size() == 3 && absl::SimpleAtoi(block[0], &tmp[0]) &&
          absl::SimpleAtoi(block[1], &tmp[1]) &&
          absl::SimpleAtoi(block[2], &tmp[2])) {
        std::copy_n(tmp, 3, kernel->mutable_block_dim()->begin());
      }
    } else if (key == "grid") {
      const std::vector<absl::string_view>& grid =
          absl::StrSplit(value_str, ',');
      uint32 tmp[3];
      if (grid.size() == 3 && absl::SimpleAtoi(grid[0], &tmp[0]) &&
          absl::SimpleAtoi(grid[1], &tmp[1]) &&
          absl::SimpleAtoi(grid[2], &tmp[2])) {
        std::copy_n(tmp, 3, kernel->mutable_grid_dim()->begin());
      }
    } else if (key == "occ_pct" && absl::SimpleAtod(value_str, &pct)) {
      kernel->set_occupancy_pct(pct);
    }
  }
}

bool IsKernelUsingTensorCore(absl::string_view kernel_name) {
  // Some examples: volta_h884gemm, volta_fp16_s884gemm,
  // turing_fp16_s1688cudnn_fp16
  bool possible_tensor_kernel = absl::StrContains(kernel_name, "884") ||
                                absl::StrContains(kernel_name, "1688") ||
                                absl::StrContains(kernel_name, "hmma") ||
                                absl::StrContains(kernel_name, "xmma");
  if (possible_tensor_kernel) {
    VLOG(3) << "Possible tensor kernel: " << kernel_name;
  }

  return (absl::StartsWith(kernel_name, "volta_i884") ||
          absl::StartsWith(kernel_name, "volta_h884") ||
          absl::StartsWith(kernel_name, "volta_s884") ||
          absl::StartsWith(kernel_name, "volta_fp16_i884") ||
          absl::StartsWith(kernel_name, "volta_fp16_h884") ||
          absl::StartsWith(kernel_name, "volta_fp16_s884") ||
          absl::StartsWith(kernel_name, "turing_i1688") ||
          absl::StartsWith(kernel_name, "turing_h1688") ||
          absl::StartsWith(kernel_name, "turing_s1688") ||
          absl::StartsWith(kernel_name, "turing_fp16_i1688") ||
          absl::StartsWith(kernel_name, "turing_fp16_h1688") ||
          absl::StartsWith(kernel_name, "turing_fp16_s1688") ||
          absl::StrContains(kernel_name, "hmma") ||
          absl::StrContains(kernel_name, "xmma"));
}

// This list is not exhaustive.
bool IsOpTensorCoreEligible(absl::string_view tf_op_name) {
  // Disable formatting to keep inline comments vertically aligned.
  // clang-format off
  return false
      // Using EndsWith to match Fused operations.
      || absl::EndsWith(tf_op_name, "Conv2D")
      || absl::EndsWith(tf_op_name, "Conv2DBackpropFilter")
      || absl::EndsWith(tf_op_name, "Conv2DBackpropInput")
      || absl::EndsWith(tf_op_name, "Conv3D")
      || absl::EndsWith(tf_op_name, "DepthwiseConv2dNative")
      || absl::EndsWith(tf_op_name, "DepthwiseConv2dNativeBackpropFilter")
      || absl::EndsWith(tf_op_name, "DepthwiseConv2dNativeBackpropInput")
      // Using Contains to match V2/V3 suffixes.
      || absl::StrContains(tf_op_name, "BatchMatMul")
      // MatMul requires exact matching.
      || absl::EndsWith(tf_op_name, "/MatMul")
      || absl::EndsWith(tf_op_name, "FusedMatMul")
      // cuDNN operations.
      || absl::EndsWith(tf_op_name, "/CudnnRNN")
      || absl::StrContains(tf_op_name, "CudnnRNNV")
      || absl::StrContains(tf_op_name, "CudnnRNNForward")
      || absl::StrContains(tf_op_name, "CudnnRNNBackprop")
      // Special cases.
      || absl::EndsWith(tf_op_name, "XlaDot")
      || absl::EndsWith(tf_op_name, "XlaDotV2");
  // clang-format on
}

bool IsEinsumTensorCoreEligible(absl::string_view equation) {
  if (equation.empty()) {
    return false;
  }
  const std::vector<absl::string_view> input_output =
      absl::StrSplit(equation, "->");
  if (input_output.size() != 2) {
    return false;
  }
  const std::vector<absl::string_view> lhs_rhs =
      absl::StrSplit(input_output[0], ',');
  return lhs_rhs.size() == 2;
}

bool KernelReportLessThanComparator::operator()(const KernelReport& lhs,
                                                const KernelReport& rhs) const {
  // Disable formatting to keep vertical alignment for better readability,
  // and make it easier to reorder columns.
  // clang-format off
  auto lhs_tuple = std::make_tuple(
      lhs.name(),
      lhs.grid_dim(0),
      lhs.grid_dim(1),
      lhs.grid_dim(2),
      lhs.block_dim(0),
      lhs.block_dim(1),
      lhs.block_dim(2),
      lhs.registers_per_thread(),
      lhs.static_shmem_bytes(),
      lhs.dynamic_shmem_bytes(),
      lhs.is_kernel_using_tensor_core(),
      lhs.is_op_tensor_core_eligible(),
      lhs.op_name());

  auto rhs_tuple = std::make_tuple(
      rhs.name(),
      rhs.grid_dim(0),
      rhs.grid_dim(1),
      rhs.grid_dim(2),
      rhs.block_dim(0),
      rhs.block_dim(1),
      rhs.block_dim(2),
      rhs.registers_per_thread(),
      rhs.static_shmem_bytes(),
      rhs.dynamic_shmem_bytes(),
      rhs.is_kernel_using_tensor_core(),
      rhs.is_op_tensor_core_eligible(),
      rhs.op_name());
  // clang-format on
  return lhs_tuple < rhs_tuple;
}

bool KernelReportEqualToComparator::operator()(const KernelReport& lhs,
                                               const KernelReport& rhs) const {
  // Disable formatting to keep vertical alignment for better readability,
  // and make it easier to reorder columns.
  // clang-format off
  // Put the most expensive string comparisons last.
  return (
      lhs.is_kernel_using_tensor_core() == rhs.is_kernel_using_tensor_core() &&
      lhs.is_op_tensor_core_eligible() == rhs.is_op_tensor_core_eligible() &&
      lhs.block_dim(0) == rhs.block_dim(0) &&
      lhs.block_dim(1) == rhs.block_dim(1) &&
      lhs.block_dim(2) == rhs.block_dim(2) &&
      lhs.grid_dim(0) == rhs.grid_dim(0) &&
      lhs.grid_dim(1) == rhs.grid_dim(1) &&
      lhs.grid_dim(2) == rhs.grid_dim(2) &&
      lhs.registers_per_thread() == rhs.registers_per_thread() &&
      lhs.static_shmem_bytes() == rhs.static_shmem_bytes() &&
      lhs.dynamic_shmem_bytes() == rhs.dynamic_shmem_bytes() &&
      lhs.name() == rhs.name() &&
      lhs.op_name() == rhs.op_name());
  // clang-format on
}

void SortAndKeepTopKDurationKernelReportsInDb(KernelStatsDb* kernel_stats_db) {
  auto comp = [](const KernelReport& lhs, const KernelReport& rhs) {
    return lhs.total_duration_ns() > rhs.total_duration_ns() ||
           (lhs.total_duration_ns() == rhs.total_duration_ns() &&
            KernelReportLessThanComparator()(lhs, rhs));
  };

  // Sort and keep at most <kMaxNumOfKernels> kernel reports.
  if (kernel_stats_db->reports_size() > kMaxNumOfKernels) {
    std::partial_sort(
        kernel_stats_db->mutable_reports()->begin(),
        kernel_stats_db->mutable_reports()->begin() + kMaxNumOfKernels,
        kernel_stats_db->mutable_reports()->end(), comp);
    kernel_stats_db->mutable_reports()->erase(
        kernel_stats_db->mutable_reports()->begin() + kMaxNumOfKernels,
        kernel_stats_db->mutable_reports()->end());
  } else {
    std::sort(kernel_stats_db->mutable_reports()->begin(),
              kernel_stats_db->mutable_reports()->end(), comp);
  }
}

void CopyTopKDurationKernelReportsToDb(const KernelReportMap& reports,
                                       KernelStatsDb* dst) {
  std::vector<std::pair<const KernelReport*, const KernelReportValue*>>
      kernels_to_sort;
  kernels_to_sort.reserve(reports.size());
  for (const auto& report_value : reports) {
    kernels_to_sort.push_back(
        std::make_pair(&report_value.first, &report_value.second));
  }

  auto comp =
      [](const std::pair<const KernelReport*, const KernelReportValue*>& lhs,
         const std::pair<const KernelReport*, const KernelReportValue*>& rhs) {
        return lhs.second->total_duration_ns > rhs.second->total_duration_ns ||
               (lhs.second->total_duration_ns ==
                    rhs.second->total_duration_ns &&
                KernelReportLessThanComparator()(*lhs.first, *rhs.first));
      };

  // Sort and copy at most <kMaxNumOfKernels> kernels to <dst>.
  if (kernels_to_sort.size() > kMaxNumOfKernels) {
    absl::c_partial_sort(kernels_to_sort,
                         kernels_to_sort.begin() + kMaxNumOfKernels, comp);
  } else {
    absl::c_sort(kernels_to_sort, comp);
  }

  int copy_size =
      std::min(kMaxNumOfKernels, static_cast<int>(kernels_to_sort.size()));
  for (int i = 0; i < copy_size; i++) {
    KernelReport* report = dst->add_reports();
    *report = *kernels_to_sort[i].first;
    const KernelReportValue& kernel_value = *kernels_to_sort[i].second;
    // Set value using KernelReportValue.
    report->set_occurrences(kernel_value.occurrences);
    report->set_min_duration_ns(kernel_value.min_duration_ns);
    report->set_max_duration_ns(kernel_value.max_duration_ns);
    report->set_total_duration_ns(kernel_value.total_duration_ns);
  }
}

void InsertOrUpdateKernelReport(const KernelReport& kernel,
                                const KernelReportValue& value,
                                KernelReportMap* dst) {
  KernelReportValue& element = (*dst)[kernel];
  if (element.occurrences == 0) {
    element = value;
  } else {
    element.total_duration_ns += value.total_duration_ns;
    element.min_duration_ns =
        std::min(element.min_duration_ns, value.min_duration_ns);
    element.max_duration_ns =
        std::max(element.max_duration_ns, value.max_duration_ns);
    element.occurrences += 1;
  }
}

void MergeKernelReports(const KernelReportMap& reports, KernelReportMap* dst) {
  for (auto& kernel_value : reports) {
    InsertOrUpdateKernelReport(kernel_value.first, kernel_value.second, dst);
  }
}

KernelStatsByOpName GroupKernelReportsByOpName(
    const KernelStatsDb& kernel_stats_db) {
  KernelStatsByOpName op_level_kernel_stats;
  for (const KernelReport& kernel_report : kernel_stats_db.reports()) {
    auto ret = op_level_kernel_stats.emplace(kernel_report.op_name(),
                                             OpLevelKernelStats());
    if (ret.second) {
      // Inserted. Add a new op in <op_level_kernel_stats>.
      OpLevelKernelStats& stats = ret.first->second;
      stats.is_op_tensor_core_eligible =
          kernel_report.is_op_tensor_core_eligible();
      stats.total_duration_ns += kernel_report.total_duration_ns();
      if (kernel_report.is_kernel_using_tensor_core()) {
        stats.tensor_core_duration_ns += kernel_report.total_duration_ns();
      }
    } else {
      // Not inserted. Aggregate kernel stats to op level.
      OpLevelKernelStats& stats = ret.first->second;
      // Verifies operations with the same name have the same TensorCore
      // eligibility.
      DCHECK_EQ(stats.is_op_tensor_core_eligible,
                kernel_report.is_op_tensor_core_eligible());
      stats.total_duration_ns += kernel_report.total_duration_ns();
      if (kernel_report.is_kernel_using_tensor_core()) {
        stats.tensor_core_duration_ns += kernel_report.total_duration_ns();
      }
    }
  }
  return op_level_kernel_stats;
}

}  // namespace profiler
}  // namespace tensorflow
<commit_msg>Update the rules to determine if a kernel uses Tensor Core.<commit_after>/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

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.
==============================================================================*/

#include "tensorflow/core/profiler/utils/kernel_stats_utils.h"

#include <algorithm>
#include <string>
#include <tuple>
#include <vector>

#include "absl/algorithm/container.h"
#include "absl/strings/match.h"
#include "absl/strings/numbers.h"
#include "absl/strings/str_split.h"
#include "absl/strings/string_view.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/profiler/protobuf/kernel_stats.pb.h"

namespace tensorflow {
namespace profiler {

namespace {

// The maximum number of Kernels displayed on Kernel Stats page.
const int kMaxNumOfKernels = 1000;

// A list of patterns to help determine if a kernel uses Tensor Core.
// A kernel uses Tensor Core if its kernel name contains any of these patterns.
// Some examples of kernel names: volta_h884gemm, turing_fp16_s1688cudnn_fp16
constexpr absl::string_view kTensorCoreKernelNamePatterns[] = {
    "16816",
    "c1688",
    "conv1x1",
    "conv2d_c1_k1",
    "dgrad_1x1_stride_2x2",
    "direct_group",
    "first_layer_wgrad_kernel",
    "h1688",
    "h884",
    "hmma",
    "i16832",
    "i8816",
    "s884",
    "s1688",
    "xmma_gemm",
    "xmma_implicit_gemm",
    "xmma_sparse_conv",
    "xmma_sparse_gemm",
    "xmma_warp_specialized_implicit_gemm"};

}  // namespace

void ParseKernelLaunchParams(absl::string_view xstat_kernel_details,
                             KernelReport* kernel) {
  const std::vector<absl::string_view> params =
      absl::StrSplit(xstat_kernel_details, absl::ByAnyChar(" \n"));

  constexpr uint32 kNumDimensions = 3;
  for (uint32 dim = 0; dim < kNumDimensions; ++dim) {
    kernel->add_block_dim(1);
    kernel->add_grid_dim(1);
  }

  // Process tokens.
  for (const auto& param : params) {
    const std::vector<absl::string_view> key_value = absl::StrSplit(param, ':');
    if (key_value.size() != 2) {
      // Unrecognized token.
      continue;
    }
    absl::string_view key = key_value[0];
    absl::string_view value_str = key_value[1];
    uint32 value = 0;
    double pct = 0.0;
    // Cases that consume a pair of tokens "key:value".
    if (key == "regs" && absl::SimpleAtoi(value_str, &value)) {
      kernel->set_registers_per_thread(value);
    } else if (key == "static_shared" && absl::SimpleAtoi(value_str, &value)) {
      kernel->set_static_shmem_bytes(value);
    } else if (key == "dynamic_shared" && absl::SimpleAtoi(value_str, &value)) {
      kernel->set_dynamic_shmem_bytes(value);
    } else if (key == "block") {
      const std::vector<absl::string_view>& block =
          absl::StrSplit(value_str, ',');
      uint32 tmp[3];
      if (block.size() == 3 && absl::SimpleAtoi(block[0], &tmp[0]) &&
          absl::SimpleAtoi(block[1], &tmp[1]) &&
          absl::SimpleAtoi(block[2], &tmp[2])) {
        std::copy_n(tmp, 3, kernel->mutable_block_dim()->begin());
      }
    } else if (key == "grid") {
      const std::vector<absl::string_view>& grid =
          absl::StrSplit(value_str, ',');
      uint32 tmp[3];
      if (grid.size() == 3 && absl::SimpleAtoi(grid[0], &tmp[0]) &&
          absl::SimpleAtoi(grid[1], &tmp[1]) &&
          absl::SimpleAtoi(grid[2], &tmp[2])) {
        std::copy_n(tmp, 3, kernel->mutable_grid_dim()->begin());
      }
    } else if (key == "occ_pct" && absl::SimpleAtod(value_str, &pct)) {
      kernel->set_occupancy_pct(pct);
    }
  }
}

bool IsKernelUsingTensorCore(absl::string_view kernel_name) {
  VLOG(1) << "kernel name: " << kernel_name;
  for (absl::string_view pattern : kTensorCoreKernelNamePatterns) {
    if (absl::StrContains(kernel_name, pattern)) {
      return true;
    }
  }
  return false;
}

// This list is not exhaustive.
bool IsOpTensorCoreEligible(absl::string_view tf_op_name) {
  // Disable formatting to keep inline comments vertically aligned.
  // clang-format off
  return false
      // Using EndsWith to match Fused operations.
      || absl::EndsWith(tf_op_name, "Conv2D")
      || absl::EndsWith(tf_op_name, "Conv2DBackpropFilter")
      || absl::EndsWith(tf_op_name, "Conv2DBackpropInput")
      || absl::EndsWith(tf_op_name, "Conv3D")
      || absl::EndsWith(tf_op_name, "DepthwiseConv2dNative")
      || absl::EndsWith(tf_op_name, "DepthwiseConv2dNativeBackpropFilter")
      || absl::EndsWith(tf_op_name, "DepthwiseConv2dNativeBackpropInput")
      // Using Contains to match V2/V3 suffixes.
      || absl::StrContains(tf_op_name, "BatchMatMul")
      // MatMul requires exact matching.
      || absl::EndsWith(tf_op_name, "/MatMul")
      || absl::EndsWith(tf_op_name, "FusedMatMul")
      // cuDNN operations.
      || absl::EndsWith(tf_op_name, "/CudnnRNN")
      || absl::StrContains(tf_op_name, "CudnnRNNV")
      || absl::StrContains(tf_op_name, "CudnnRNNForward")
      || absl::StrContains(tf_op_name, "CudnnRNNBackprop")
      // Special cases.
      || absl::EndsWith(tf_op_name, "XlaDot")
      || absl::EndsWith(tf_op_name, "XlaDotV2");
  // clang-format on
}

bool IsEinsumTensorCoreEligible(absl::string_view equation) {
  if (equation.empty()) {
    return false;
  }
  const std::vector<absl::string_view> input_output =
      absl::StrSplit(equation, "->");
  if (input_output.size() != 2) {
    return false;
  }
  const std::vector<absl::string_view> lhs_rhs =
      absl::StrSplit(input_output[0], ',');
  return lhs_rhs.size() == 2;
}

bool KernelReportLessThanComparator::operator()(const KernelReport& lhs,
                                                const KernelReport& rhs) const {
  // Disable formatting to keep vertical alignment for better readability,
  // and make it easier to reorder columns.
  // clang-format off
  auto lhs_tuple = std::make_tuple(
      lhs.name(),
      lhs.grid_dim(0),
      lhs.grid_dim(1),
      lhs.grid_dim(2),
      lhs.block_dim(0),
      lhs.block_dim(1),
      lhs.block_dim(2),
      lhs.registers_per_thread(),
      lhs.static_shmem_bytes(),
      lhs.dynamic_shmem_bytes(),
      lhs.is_kernel_using_tensor_core(),
      lhs.is_op_tensor_core_eligible(),
      lhs.op_name());

  auto rhs_tuple = std::make_tuple(
      rhs.name(),
      rhs.grid_dim(0),
      rhs.grid_dim(1),
      rhs.grid_dim(2),
      rhs.block_dim(0),
      rhs.block_dim(1),
      rhs.block_dim(2),
      rhs.registers_per_thread(),
      rhs.static_shmem_bytes(),
      rhs.dynamic_shmem_bytes(),
      rhs.is_kernel_using_tensor_core(),
      rhs.is_op_tensor_core_eligible(),
      rhs.op_name());
  // clang-format on
  return lhs_tuple < rhs_tuple;
}

bool KernelReportEqualToComparator::operator()(const KernelReport& lhs,
                                               const KernelReport& rhs) const {
  // Disable formatting to keep vertical alignment for better readability,
  // and make it easier to reorder columns.
  // clang-format off
  // Put the most expensive string comparisons last.
  return (
      lhs.is_kernel_using_tensor_core() == rhs.is_kernel_using_tensor_core() &&
      lhs.is_op_tensor_core_eligible() == rhs.is_op_tensor_core_eligible() &&
      lhs.block_dim(0) == rhs.block_dim(0) &&
      lhs.block_dim(1) == rhs.block_dim(1) &&
      lhs.block_dim(2) == rhs.block_dim(2) &&
      lhs.grid_dim(0) == rhs.grid_dim(0) &&
      lhs.grid_dim(1) == rhs.grid_dim(1) &&
      lhs.grid_dim(2) == rhs.grid_dim(2) &&
      lhs.registers_per_thread() == rhs.registers_per_thread() &&
      lhs.static_shmem_bytes() == rhs.static_shmem_bytes() &&
      lhs.dynamic_shmem_bytes() == rhs.dynamic_shmem_bytes() &&
      lhs.name() == rhs.name() &&
      lhs.op_name() == rhs.op_name());
  // clang-format on
}

void SortAndKeepTopKDurationKernelReportsInDb(KernelStatsDb* kernel_stats_db) {
  auto comp = [](const KernelReport& lhs, const KernelReport& rhs) {
    return lhs.total_duration_ns() > rhs.total_duration_ns() ||
           (lhs.total_duration_ns() == rhs.total_duration_ns() &&
            KernelReportLessThanComparator()(lhs, rhs));
  };

  // Sort and keep at most <kMaxNumOfKernels> kernel reports.
  if (kernel_stats_db->reports_size() > kMaxNumOfKernels) {
    std::partial_sort(
        kernel_stats_db->mutable_reports()->begin(),
        kernel_stats_db->mutable_reports()->begin() + kMaxNumOfKernels,
        kernel_stats_db->mutable_reports()->end(), comp);
    kernel_stats_db->mutable_reports()->erase(
        kernel_stats_db->mutable_reports()->begin() + kMaxNumOfKernels,
        kernel_stats_db->mutable_reports()->end());
  } else {
    std::sort(kernel_stats_db->mutable_reports()->begin(),
              kernel_stats_db->mutable_reports()->end(), comp);
  }
}

void CopyTopKDurationKernelReportsToDb(const KernelReportMap& reports,
                                       KernelStatsDb* dst) {
  std::vector<std::pair<const KernelReport*, const KernelReportValue*>>
      kernels_to_sort;
  kernels_to_sort.reserve(reports.size());
  for (const auto& report_value : reports) {
    kernels_to_sort.push_back(
        std::make_pair(&report_value.first, &report_value.second));
  }

  auto comp =
      [](const std::pair<const KernelReport*, const KernelReportValue*>& lhs,
         const std::pair<const KernelReport*, const KernelReportValue*>& rhs) {
        return lhs.second->total_duration_ns > rhs.second->total_duration_ns ||
               (lhs.second->total_duration_ns ==
                    rhs.second->total_duration_ns &&
                KernelReportLessThanComparator()(*lhs.first, *rhs.first));
      };

  // Sort and copy at most <kMaxNumOfKernels> kernels to <dst>.
  if (kernels_to_sort.size() > kMaxNumOfKernels) {
    absl::c_partial_sort(kernels_to_sort,
                         kernels_to_sort.begin() + kMaxNumOfKernels, comp);
  } else {
    absl::c_sort(kernels_to_sort, comp);
  }

  int copy_size =
      std::min(kMaxNumOfKernels, static_cast<int>(kernels_to_sort.size()));
  for (int i = 0; i < copy_size; i++) {
    KernelReport* report = dst->add_reports();
    *report = *kernels_to_sort[i].first;
    const KernelReportValue& kernel_value = *kernels_to_sort[i].second;
    // Set value using KernelReportValue.
    report->set_occurrences(kernel_value.occurrences);
    report->set_min_duration_ns(kernel_value.min_duration_ns);
    report->set_max_duration_ns(kernel_value.max_duration_ns);
    report->set_total_duration_ns(kernel_value.total_duration_ns);
  }
}

void InsertOrUpdateKernelReport(const KernelReport& kernel,
                                const KernelReportValue& value,
                                KernelReportMap* dst) {
  KernelReportValue& element = (*dst)[kernel];
  if (element.occurrences == 0) {
    element = value;
  } else {
    element.total_duration_ns += value.total_duration_ns;
    element.min_duration_ns =
        std::min(element.min_duration_ns, value.min_duration_ns);
    element.max_duration_ns =
        std::max(element.max_duration_ns, value.max_duration_ns);
    element.occurrences += 1;
  }
}

void MergeKernelReports(const KernelReportMap& reports, KernelReportMap* dst) {
  for (auto& kernel_value : reports) {
    InsertOrUpdateKernelReport(kernel_value.first, kernel_value.second, dst);
  }
}

KernelStatsByOpName GroupKernelReportsByOpName(
    const KernelStatsDb& kernel_stats_db) {
  KernelStatsByOpName op_level_kernel_stats;
  for (const KernelReport& kernel_report : kernel_stats_db.reports()) {
    auto ret = op_level_kernel_stats.emplace(kernel_report.op_name(),
                                             OpLevelKernelStats());
    if (ret.second) {
      // Inserted. Add a new op in <op_level_kernel_stats>.
      OpLevelKernelStats& stats = ret.first->second;
      stats.is_op_tensor_core_eligible =
          kernel_report.is_op_tensor_core_eligible();
      stats.total_duration_ns += kernel_report.total_duration_ns();
      if (kernel_report.is_kernel_using_tensor_core()) {
        stats.tensor_core_duration_ns += kernel_report.total_duration_ns();
      }
    } else {
      // Not inserted. Aggregate kernel stats to op level.
      OpLevelKernelStats& stats = ret.first->second;
      // Verifies operations with the same name have the same TensorCore
      // eligibility.
      DCHECK_EQ(stats.is_op_tensor_core_eligible,
                kernel_report.is_op_tensor_core_eligible());
      stats.total_duration_ns += kernel_report.total_duration_ns();
      if (kernel_report.is_kernel_using_tensor_core()) {
        stats.tensor_core_duration_ns += kernel_report.total_duration_ns();
      }
    }
  }
  return op_level_kernel_stats;
}

}  // namespace profiler
}  // namespace tensorflow
<|endoftext|>


<commit_before>e378d98c-2747-11e6-8d69-e0f84713e7b8<commit_msg>TODO: Add a beter commit message<commit_after>e3a37b91-2747-11e6-88f5-e0f84713e7b8<|endoftext|>


<commit_before>#include <iostream>
#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <vector>
#include <cstring>
#include <sys/unistd.h>
#include <cstdlib>

using namespace std;

class Input
{
    protected:
        string input;

    public:
        Input()
        {
            input = "";
        }

        Input(string s)
        {
            input = s;
        }

        void recieveinput(int &check)
        {
            cout << "$: ";
            getline(cin, input);
            check = 0;
        }

        string get_string()
        {
            return input;
        }
};

class Command : public Input
{
    protected:
        vector<string> v;

    public:
        Command()
        {
            for(unsigned int i = 0; i < v.size(); ++i)
            {
                v.at(i) = "";
            }
        }

        vector<string> get_vector()
        {
            return v;
        }

        void parse_string(string s)
        {
            unsigned int index = 0;
            string temp = "";
            unsigned int i = 0;
            unsigned int space = i;

            for(i = 0; i < s.length(); ++i)
            {
                for(i = space; i < s.length(); ++i)
                {
                    if(s.at(i) != ' ' && s.at(i) != '\0')
                    {
                        temp += s.at(i);
                    }
                    else if(s.at(i) == ' ' || s.at(i) == '\0')
                    {
                        break;
                    }

                }
                v.push_back(temp);
                temp = "";
                ++index;
                space = i + 1;
            }
        }

        void exc_command(string command, int& passed)
        {
            passed = 1;
            char* arg[256];
            string tempA = "";
            string tempB = "";
            string tempC = "";
            string tempD = "";
            int counter = 0;

            for(unsigned int i = 0; i < command.size(); i++)
            {
                if((command.at(i) != ' ' ) && (counter == 0))
                {
                    tempA += command.at(i);
                }
                else if((command.at(i) != ' ') && (counter == 1))
                {
                    tempB += command.at(i);
                }
                else if((command.at(i) != ' ') && (counter == 2))
                {
                    tempC += command.at(i);
                }
                else if((command.at(i) != ' ') && (counter == 3))
                {
                    tempD += command.at(i);
                }
                if(command.at(i) == ' ')
                {
                    counter++;
                }
            }
            if(tempB.size() == 0)
            {
                arg[0] = (char*)tempA.c_str();
                arg[1] = NULL;
            }
            else if(tempC.size() == 0)
            {
                arg[0] = (char*)tempA.c_str();
                arg[1] = (char*)tempB.c_str();
                arg[2] = NULL;
            }
            else if(tempD.size() == 0)
            {
                arg[0] = (char*)tempA.c_str();
                arg[1] = (char*)tempB.c_str();
                arg[2] = (char*)tempC.c_str();
                arg[4] = NULL;
            }
            else
            {
                arg[0] = (char*)tempA.c_str();
                arg[1] = (char*)tempB.c_str();
                arg[2] = (char*)tempC.c_str();
                arg[4] = (char*)tempD.c_str();
                arg[5] = NULL;
            }

            pid_t pid = fork();
            if(pid == 0)
            {
                if(execvp(arg[0],arg) == -1)
                {
                    perror("exec");
                    passed = 0;
                }
            }
            if(pid > 1)
            {
                if(wait(0) == -1)
                {
                    perror("wait");
                    passed = 0;
                }
            }
        }
};

int main()
{
    int check = 0;
    vector<char> vc1;
    vector<char> vc2;
    vector<string> vs1;
    vector<string> vs2;

    Input in;
    Command vc;

    in.recieveinput(check);
    vc.parse_string(in.get_string());

    unsigned int index = 0;
    while(in.get_string() != "exit")
    {
        for(unsigned int i = 0; i < in.get_string().size(); i++)
        {
            if((in.get_string().at(i) == ';') || (in.get_string().at(i) == '&') || (in.get_string().at(i) == '|'))
            {
                check = 1;
                vc1.push_back(in.get_string().at(i));
            }
            ++index;
        }

        if(check == 0)
        {
            int cmdcount = 1;
            vc.exc_command(in.get_string(), cmdcount);
        }

        else if(check == 1)
        {
            string tmp = "";

            char *inp;

            inp = new char[in.get_string().length()];
            for(unsigned int i = 0; i < in.get_string().length(); i++)
            {
                inp[i] = in.get_string()[i];
            }
            inp[in.get_string().length()] = '\0';

            char *pnt;
            pnt = strtok(inp, ";&&||");
            while(pnt != NULL)
            {
                int i = 0;
                while((pnt[i] != '\0') && (pnt[i] != '#'))
                {
                    tmp += pnt[i];
                    i++;
                }
                if(tmp.at(0) == ' ')
                {
                    string tmpA = tmp;
                    tmp.clear();
                    for(unsigned int i = 0; i < tmpA.size(); i++)
                    {
                        if(i == 0)
                        {}
                        else
                        {
                            tmp += tmpA.at(i);
                        }
                    }
                }
                vs1.push_back(tmp);
                pnt = strtok(NULL,";&&||");
                tmp.clear();
            }
            vs2 = vs1;
            vc2 = vc1;
            vc1.clear();
            vs1.clear();
            for(int i = vs2.size() - 1; i > -1; i--)
            {
                vs1.push_back(vs2.at(i));
            }
            for(int i = vc2.size() - 1; i > -1; i--)
            {
                vc1.push_back(vc2.at(i));
            }

            int cmdcnt = 0;
            int skp = 0;

            char track = vc1.at(vc1.size() - 1);
            string track2 = vs1.at(vs1.size() - 1);

            Command vect;
            vect.parse_string(track2);

            vect.exc_command(track2, cmdcnt);
            vs1.pop_back();
            do
            {
                if(track == ';')
                {
                    vc1.pop_back();
                    skp = 1;
                }
                if((track == '&') && (cmdcnt == 1) && (skp == 0))
                {
                    vc1.pop_back();
                    vc1.pop_back();

                    if((vs1.size() != 0) && (vc1.size() != 0))
                    {
                        track = vc1.at(vc1.size() - 1);
                        track2 = vs1.at(vs1.size() - 1);

                        vect.get_vector().clear();
                        vect.parse_string(track2);
                        vect.exc_command(track2, cmdcnt);
                        vs1.pop_back();
                    }
                    skp = 1;
                }
                else if((track == '&') && (cmdcnt == 0) && (skp == 0))
                {
                    vc1.pop_back();
                    vc1.pop_back();
                    vs1.pop_back();
                    skp = 1;
                }
                if((track == '|') && (cmdcnt == 0) && (skp == 0))
                {
                    vc1.pop_back();
                    vc1.pop_back();
                    if(vs1.size() != 0)
                    {
                        track2 = vs1.at(vs1.size() - 1);
                        vect.get_vector().clear();
                        vect.parse_string(track2);
                        vect.exc_command(track2, cmdcnt);
                        vs1.pop_back();
                    }
                    skp = 1;
                }
                else if((track == '|') && (cmdcnt == 1 ) && (skp = 0))
                {
                        vc1.pop_back();
                        vc1.pop_back();
                        vs1.pop_back();
                        skp = 1;
                }
                skp = 0;
                if(vc1.size() != 0)
                {
                    track = vc1.at(vc1.size() - 1);
                }
                else
                {
                    track = '0';
                }
                if(vs1.size() != 0)
                {
                    track2 = vs1.at(vs1.size() - 1);
                }
                else
                {
                    track2 = "";
                }
                if((track == ';') && (track2.size() > 1))
                {
                    vect.get_vector().clear();
                    vect.parse_string(track2);
                    vect.exc_command(track2, cmdcnt);
                    vs1.pop_back();
                }
                else if((track == '|') && (track2.size() > 1))
                {
                    if(cmdcnt == 0)
                    {
                        vect.get_vector().clear();
                        vect.parse_string(track2);
                        vect.exc_command(track2, cmdcnt);
                        vs1.pop_back();
                    }
                }
                else if((track == '&') && (track2.size() > 1))
                {
                    if(cmdcnt == 1)
                    {
                        vect.parse_string(track2);
                        vect.exc_command(track2,cmdcnt);
                        vs1.pop_back();
                    }
                }
                else
                {
                    if(track2.size() > 1)
                    {
                        vect.get_vector().clear();
                        vect.parse_string(track2);
                        vect.exc_command(track2, cmdcnt);
                        vs1.pop_back();
                    }
                }
            }
            while((vc1.size() != 0) && (vs1.size() != 0));
        }
        vs1.clear();
        vs2.clear();
        vc1.clear();
        vc2.clear();
        in.get_string().clear();
        in.recieveinput(check);
    }

    return 0;
}<commit_msg>final changes<commit_after>#include <iostream>
#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <vector>
#include <cstring>
#include <sys/unistd.h>
#include <cstdlib>

using namespace std;

class Input
{
    protected:
        string input;

    public:
        Input()
        {
            input = "";
        }

        Input(string s)
        {
            input = s;
        }

        void recieveinput(int &check)
        {
            cout << "$: ";
            getline(cin, input);
            check = 0;
        }

        string get_string()
        {
            return input;
        }
};

class Command : public Input
{
    protected:
        vector<string> v;

    public:
        Command()
        {
            for(unsigned int i = 0; i < v.size(); ++i)
            {
                v.at(i) = "";
            }
        }

        vector<string> get_vector()
        {
            return v;
        }

        void parse_string(string s)
        {
            unsigned int index = 0;
            string temp = "";
            unsigned int i = 0;
            unsigned int space = i;

            for(i = 0; i < s.length(); ++i)
            {
                for(i = space; i < s.length(); ++i)
                {
                    if(s.at(i) != ' ' && s.at(i) != '\0')
                    {
                        temp += s.at(i);
                    }
                    else if(s.at(i) == ' ' || s.at(i) == '\0')
                    {
                        break;
                    }

                }
                v.push_back(temp);
                temp = "";
                ++index;
                space = i + 1;
            }
        }

        void exc_command(string command, int& passed)
        {
            passed = 1;
            char* arg[256];
            string tempA = "";
            string tempB = "";
            string tempC = "";
            string tempD = "";
            int counter = 0;

            for(unsigned int i = 0; i < command.size(); i++)
            {
                if((command.at(i) != ' ' ) && (counter == 0))
                {
                    tempA += command.at(i);
                }
                else if((command.at(i) != ' ') && (counter == 1))
                {
                    tempB += command.at(i);
                }
                else if((command.at(i) != ' ') && (counter == 2))
                {
                    tempC += command.at(i);
                }
                else if((command.at(i) != ' ') && (counter == 3))
                {
                    tempD += command.at(i);
                }
                if(command.at(i) == ' ')
                {
                    counter++;
                }
            }
            if(tempB.size() == 0)
            {
                arg[0] = (char*)tempA.c_str();
                arg[1] = NULL;
            }
            else if(tempC.size() == 0)
            {
                arg[0] = (char*)tempA.c_str();
                arg[1] = (char*)tempB.c_str();
                arg[2] = NULL;
            }
            else if(tempD.size() == 0)
            {
                arg[0] = (char*)tempA.c_str();
                arg[1] = (char*)tempB.c_str();
                arg[2] = (char*)tempC.c_str();
                arg[4] = NULL;
            }
            else
            {
                arg[0] = (char*)tempA.c_str();
                arg[1] = (char*)tempB.c_str();
                arg[2] = (char*)tempC.c_str();
                arg[4] = (char*)tempD.c_str();
                arg[5] = NULL;
            }

            pid_t pid = fork();
            if(pid == 0)
            {
                if(execvp(arg[0],arg) == -1)
                {
                    perror("exec");
                    passed = 0;
                }
            }
            if(pid > 1)
            {
                if(wait(0) == -1)
                {
                    perror("wait");
                    passed = 0;
                }
            }
        }
};

int main()
{
    int check = 0;
    vector<char> vc1;
    vector<char> vc2;
    vector<string> vs1;
    vector<string> vs2;

    Input in;
    Command vc;

    in.recieveinput(check);
    vc.parse_string(in.get_string());

    unsigned int index = 0;
    while(in.get_string() != "exit")
    {
        for(unsigned int i = 0; i < in.get_string().size(); i++)
        {
            if((in.get_string().at(i) == ';') || (in.get_string().at(i) == '&') || (in.get_string().at(i) == '|'))
            {
                check = 1;
                vc1.push_back(in.get_string().at(i));
            }
            ++index;
        }

        if(check == 0)
        {
            int cmdcount = 1;
            vc.exc_command(in.get_string(), cmdcount);
        }

        else if(check == 1)
        {
            string tmp = "";

            char *inp;

            inp = new char[in.get_string().length()];
            for(unsigned int i = 0; i < in.get_string().length(); i++)
            {
                inp[i] = in.get_string()[i];
            }
            inp[in.get_string().length()] = '\0';

            char *pnt;
            pnt = strtok(inp, ";&&||");
            while(pnt != NULL)
            {
                int i = 0;
                while((pnt[i] != '\0') && (pnt[i] != '#'))
                {
                    tmp += pnt[i];
                    i++;
                }
                if(tmp.at(0) == ' ')
                {
                    string tmpA = tmp;
                    tmp.clear();
                    for(unsigned int i = 0; i < tmpA.size(); i++)
                    {
                        if(i == 0)
                        {}
                        else
                        {
                            tmp += tmpA.at(i);
                        }
                    }
                }
                vs1.push_back(tmp);
                pnt = strtok(NULL,";&&||");
                tmp.clear();
            }
            vs2 = vs1;
            vc2 = vc1;
            vc1.clear();
            vs1.clear();
            for(int i = vs2.size() - 1; i > -1; i--)
            {
                vs1.push_back(vs2.at(i));
            }
            for(int i = vc2.size() - 1; i > -1; i--)
            {
                vc1.push_back(vc2.at(i));
            }

            int cmdcnt = 0;
            int skp = 0;

            char track = vc1.at(vc1.size() - 1);
            string track2 = vs1.at(vs1.size() - 1);

            Command vect;
            vect.parse_string(track2);

            vect.exc_command(track2, cmdcnt);
            vs1.pop_back();
            do
            {
                if(track == ';')
                {
                    vc1.pop_back();
                    skp = 1;
                }
                if((track == '&') && (cmdcnt == 1) && (skp == 0))
                {
                    vc1.pop_back();
                    vc1.pop_back();

                    if((vs1.size() != 0) && (vc1.size() != 0))
                    {
                        track = vc1.at(vc1.size() - 1);
                        track2 = vs1.at(vs1.size() - 1);

                        vect.get_vector().clear();
                        vect.parse_string(track2);
                        vect.exc_command(track2, cmdcnt);
                        vs1.pop_back();
                    }
                    skp = 1;
                }
                else if((track == '&') && (cmdcnt == 0) && (skp == 0))
                {
                    vc1.pop_back();
                    vc1.pop_back();
                    vs1.pop_back();
                    skp = 1;
                }
                if((track == '|') && (cmdcnt == 0) && (skp == 0))
                {
                    vc1.pop_back();
                    vc1.pop_back();
                    if(vs1.size() != 0)
                    {
                        track2 = vs1.at(vs1.size() - 1);
                        vect.get_vector().clear();
                        vect.parse_string(track2);
                        vect.exc_command(track2, cmdcnt);
                        vs1.pop_back();
                    }
                    skp = 1;
                }
                else if((track == '|') && (cmdcnt == 1 ) && (skp = 0))
                {
                        vc1.pop_back();
                        vc1.pop_back();
                        vs1.pop_back();
                        skp = 1;
                }
                skp = 0;
                if(vc1.size() != 0)
                {
                    track = vc1.at(vc1.size() - 1);
                }
                else
                {
                    track = '0';
                }
                if(vs1.size() != 0)
                {
                    track2 = vs1.at(vs1.size() - 1);
                }
                else
                {
                    track2 = "";
                }
                if((track == ';') && (track2.size() > 1))
                {
                    vect.get_vector().clear();
                    vect.parse_string(track2);
                    vect.exc_command(track2, cmdcnt);
                    vs1.pop_back();
                }
                else if((track == '|') && (track2.size() > 1))
                {
                    if(cmdcnt == 0)
                    {
                        vect.get_vector().clear();
                        vect.parse_string(track2);
                        vect.exc_command(track2, cmdcnt);
                        vs1.pop_back();
                    }
                }
                else if((track == '&') && (track2.size() > 1))
                {
                    if(cmdcnt == 1)
                    {
                        vect.parse_string(track2);
                        vect.exc_command(track2,cmdcnt);
                        vs1.pop_back();
                    }
                }
                else
                {
                    if(track2.size() > 1)
                    {
                        vect.get_vector().clear();
                        vect.parse_string(track2);
                        vect.exc_command(track2, cmdcnt);
                        vs1.pop_back();
                    }
                }
            }
            while((vc1.size() != 0) && (vs1.size() != 0));
        }
        vs1.clear();
        vs2.clear();
        vc1.clear();
        vc2.clear();
        in.get_string().clear();
        in.recieveinput(check);
    }

    return 0;
}
<|endoftext|>


<commit_before>#include "yael/network/TcpSocket.h"

#include <sstream>
#include <fcntl.h>
#include <sys/socket.h>
#include <sys/poll.h>
#include <sys/types.h>
#include <csignal>
#include <arpa/inet.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <iostream>
#include <cerrno>
#include <stdexcept>
#include <cstring>
#include <unistd.h>
#include <cassert>
#include <glog/logging.h>

#include "DatagramMessageSlicer.h"
#include "StreamMessageSlicer.h"

using namespace std;

namespace yael::network
{

constexpr int TRUE_FLAG = 1;

TcpSocket::TcpSocket(MessageMode mode, size_t max_send_queue_size)
    : m_port(0), m_is_ipv6(false), m_fd(-1), m_max_send_queue_size(max_send_queue_size)
{
    if(mode == MessageMode::Datagram)
    {
        m_slicer = std::make_unique<DatagramMessageSlicer>();
    }
    else if(mode == MessageMode::Stream)
    {
        m_slicer = std::make_unique<StreamMessageSlicer>();
    }
    else
    {
        throw std::runtime_error("Invalid message mode");
    }
}

TcpSocket::TcpSocket(MessageMode mode, int fd, size_t max_send_queue_size)
    : m_port(0), m_is_ipv6(false), m_fd(fd), m_max_send_queue_size(max_send_queue_size)
{
    if(mode == MessageMode::Datagram)
    {
        m_slicer = std::make_unique<DatagramMessageSlicer>();
    }
    else if(mode == MessageMode::Stream)
    {
        m_slicer = std::make_unique<StreamMessageSlicer>();
    }
    else
    {
        throw std::runtime_error("Invalid message mode");
    }

    uint32_t flags = fcntl(m_fd, F_GETFL, 0);
    flags = flags | O_NONBLOCK;
    fcntl(m_fd, F_SETFL, flags);

    update_port_number();
    calculate_remote_address();

    m_state = State::Connected;
}

TcpSocket::~TcpSocket()
{
    TcpSocket::close(true);
}

bool TcpSocket::has_messages() const
{
    return m_slicer->has_messages();
}

bool TcpSocket::create_fd()
{
    if(m_is_ipv6)
    {
        m_fd = socket(AF_INET6, SOCK_STREAM, IPPROTO_TCP);
    }
    else
    {
        m_fd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    }

    if(!is_valid())
    {
        return false;
    }

    ::setsockopt(m_fd, IPPROTO_TCP, TCP_NODELAY, reinterpret_cast<const char*>(&TRUE_FLAG), sizeof(TRUE_FLAG));

    return true;
}

void TcpSocket::update_port_number()
{
    if(m_is_ipv6)
    {
        struct sockaddr_in6 addr;
        socklen_t addrlen = sizeof(addr);
        getsockname(m_fd, reinterpret_cast<sockaddr*>(&addr), &addrlen);
        m_port = htons(addr.sin6_port);
    }
    else
    {
        struct sockaddr_in addr;
        socklen_t addrlen = sizeof(addr);
        getsockname(m_fd, reinterpret_cast<sockaddr*>(&addr), &addrlen);
        m_port = htons(addr.sin_port);
    }
}

bool TcpSocket::bind_socket(const Address& address)
{
    m_is_ipv6 = address.IPv6;

    if(!create_fd())
    {
        return false;
    }

    // Reuse address so we can quickly recover from crashes
    ::setsockopt(m_fd, SOL_SOCKET, SO_REUSEADDR, &TRUE_FLAG, sizeof(TRUE_FLAG));

    if(m_is_ipv6)
    {
        sockaddr_in6 sock_addr;
        address.get_sock_address6(sock_addr);

        if(::bind(m_fd, reinterpret_cast<sockaddr*>(&sock_addr), sizeof(sock_addr)) != 0)
        {
            return false;
        }
    }
    else
    {
        sockaddr_in sock_addr;
        address.get_sock_address(sock_addr);

        if(::bind(m_fd, reinterpret_cast<sockaddr*>(&sock_addr), sizeof(sock_addr)) != 0)
        {
            return false;
        }
   }

    update_port_number();
    return true;
}

bool TcpSocket::listen(const Address& address, uint32_t backlog)
{
    if(!bind_socket(address))
    {
        throw socket_error("Failed to bind socket!");
    }

    uint32_t flags = fcntl(m_fd, F_GETFL, 0);
    flags = flags | O_NONBLOCK;
    fcntl(m_fd, F_SETFL, flags);

    if(::listen(m_fd, backlog) == 0)
    {
        m_state = State::Listening;
        return true;
    }
    else
    {
        return false;
    }
}

uint16_t TcpSocket::port() const
{
    if(!is_valid())
    {
        throw socket_error("Cannot get port of non-existent socket");
    }

    return m_port;
}

bool TcpSocket::connect(const Address& address, const std::string& name)
{
    if(address.PortNumber == 0)
    {
        throw std::invalid_argument("Need to specify a port number");
    }

    if(name.empty())
    {
        m_is_ipv6 = address.IPv6;
        if(!create_fd())
        {
            throw socket_error(strerror(errno));
        }
    }
    else
    {
        Address my_addr = resolve_URL(name, ANY_PORT, address.IPv6);

        if(!bind_socket(my_addr))
        {
            throw socket_error(strerror(errno));
        }
    }

    // Set it blocking just for connect
    uint32_t flags = fcntl(m_fd, F_GETFL, 0);
    flags = flags & ~O_NONBLOCK;
    fcntl(m_fd, F_SETFL, flags);

    if(m_is_ipv6)
    {
        sockaddr_in6 sock_addr;
        address.get_sock_address6(sock_addr);

        if(::connect(m_fd, reinterpret_cast<const sockaddr*>(&sock_addr), sizeof(sock_addr)) != 0)
        {
            close();
            return false;
        }
    }
    else
    {
        sockaddr_in sock_addr;
        address.get_sock_address(sock_addr);

        if(::connect(m_fd, reinterpret_cast<const sockaddr*>(&sock_addr), sizeof(sock_addr)) != 0)
        {
            close();
            return false;
        }
    }

    flags = fcntl(m_fd, F_GETFL, 0);
    flags = flags | O_NONBLOCK;
    fcntl(m_fd, F_SETFL, flags);

    calculate_remote_address();
    update_port_number();

    m_state = State::Connected;
    return true;
}

int32_t TcpSocket::internal_accept()
{
    int32_t s = 0;

    if(m_is_ipv6)
    {
        sockaddr_in6 sock_addr;
        socklen_t len = sizeof(sock_addr);
        s = ::accept(m_fd, reinterpret_cast<sockaddr*>(&sock_addr), &len);
    }
    else
    {
        sockaddr_in sock_addr;
        socklen_t len = sizeof(sock_addr);
        s = ::accept(m_fd, reinterpret_cast<sockaddr*>(&sock_addr), &len);
    }

    if(s < 0 && errno != EWOULDBLOCK)
    {
        close();
        std::string str = "Failed to accept new connection; ";
        str += strerror(errno);
        throw socket_error(str);
    }

    return s;
}

std::vector<std::unique_ptr<Socket>> TcpSocket::accept()
{
    if(!is_listening())
    {
        throw socket_error("Cannot accept on connected TcpTcpSocket");
    }

    std::vector<std::unique_ptr<Socket>> res;

    while(true)
    {
        auto fd = internal_accept();

        if(fd >= 0)
        {
            auto ptr = new TcpSocket(m_slicer->type(), fd, m_max_send_queue_size);
            res.emplace_back(std::unique_ptr<Socket>(ptr));
        }
        else
        {
            return res;
        }
    }
}

const Address& TcpSocket::get_remote_address() const
{
    return m_remote_address;
}

void TcpSocket::calculate_remote_address()
{
    char ipstring[16];
    sockaddr_in sin;
    uint32_t len = sizeof(sin);

    if( getpeername(m_fd, reinterpret_cast<sockaddr*>(&sin), &len) == -1)
    {
        m_remote_address.reset();
    }

    inet_ntop( AF_INET, dynamic_cast<in_addr*>(&sin.sin_addr), &ipstring[0], 16);

    uint16_t port = htons(sin.sin_port);

    m_remote_address.IP = &ipstring[0];
    m_remote_address.PortNumber = port;
}

bool TcpSocket::close(bool fast)
{
    if(m_state == State::Connected && !fast)
    {
        m_state = State::Shutdown;
        int i = ::shutdown(m_fd, SHUT_RD | SHUT_WR);
        (void)i; //unused

        return false;
    }

    if(m_fd > 0)
    {
        m_state = State::Closed;
        int i = ::close(m_fd);
        (void)i; //unused
        m_fd = -1;

        m_slicer->buffer().reset();
    }
    else
    {
        if(!(m_state == State::Closed || m_state == State::Unknown))
        {
            LOG(FATAL) << "Invalid state";
        }

        //no-op
    }

    // threads might wait for send queue to be empty
    {
        std::unique_lock lock(m_send_mutex);
        m_send_queue_cond.notify_all();
    }

    return true;
}

void TcpSocket::pull_messages()
{
    auto &buffer = m_slicer->buffer();

    if(!buffer.is_valid())
    {
        bool res = receive_data(buffer);
        if(!res)
        {
            return;
        }
    }

    try
    {
        m_slicer->process_buffer();
    }
    catch(std::exception &e)
    {
        // ignore
    }

    // read rest of buffer
    // always pull more until we get EAGAIN
    pull_messages();
}

bool TcpSocket::receive_data(buffer_t &buffer)
{
    if(!is_valid())
    {
        return false;
    }

    if(buffer.is_valid())
    {
        throw std::runtime_error("TcpSocket::receive_data failed: Still have data queued up in buffer");
    }

    memset(&buffer.data[0], 0, yael::network::buffer_t::MAX_SIZE);
    auto x = ::recv(m_fd, buffer.data, yael::network::buffer_t::MAX_SIZE, 0);

    // Now act accordingly
    // > 0 -> data
    // = 0 -> disconnect
    // < 0 -> error/block
    if(x > 0)
    {
        buffer.size = x;
        buffer.position = 0;

        return true;
    }
    else if(x == 0)
    {
        close(true);
        return false;
    }
    else
    {
        const int e = errno;

        switch(e)
        {
        case EAGAIN:
            break;
        case ECONNRESET:
            close(true);
            break;
        default:
        {
            std::string str = "Failed to receive data; ";
            str += strerror(errno);

            // First close socket and then throw the error!
            close(true);
            throw socket_error(str);
        }
        }

        return false;
    }
}

std::optional<TcpSocket::message_in_t> TcpSocket::receive()
{
    pull_messages();

    if(m_slicer->has_messages())
    {
        message_in_t msg;
        auto res = m_slicer->get_message(msg);
        if(!res)
        {
            throw socket_error("failed to get message");
        }

        return { msg };
    }
    else
    {
        return {};
    }
}

bool TcpSocket::send(std::unique_ptr<uint8_t[]> &&data, uint32_t len)
{
    if(len <= 0)
    {
        throw socket_error("Message size has to be > 0");
    }

    m_slicer->prepare_message(data, len);

    auto msg_out = message_out_internal_t(std::move(data), len);

    {
        std::unique_lock lock(m_send_mutex);

        if(m_send_queue_size >= m_max_send_queue_size)
        {
            throw send_queue_full();
        }

        m_send_queue_size += msg_out.length;
        m_send_queue.emplace_back(std::move(msg_out));
    }

    return do_send();
}

void TcpSocket::wait_send_queue_empty()
{
    std::unique_lock lock(m_send_mutex);

    while(m_send_queue_size > 0 && is_valid())
    {
        m_send_queue_cond.wait(lock);
    }
}

bool TcpSocket::do_send()
{
    std::unique_lock lock(m_send_mutex);

    while(true)
    {
        auto it = m_send_queue.begin();

        if(it == m_send_queue.end())
        {
            // we sent everything!
            return false;
        }

        if(!is_valid())
        {
            throw socket_error("Socket is closed");
        }

        auto &message = *it;

        const uint32_t length = message.length;
        const uint8_t *rdata = message.data.get();

        while(message.sent_pos < length)
        {
            auto s = ::write(m_fd, rdata + message.sent_pos, length - message.sent_pos);

            if(s > 0)
            {
                message.sent_pos += s;
            }
            else if(s == 0)
            {
                LOG(WARNING) << "Connection lost during send: Message may only be sent partially";
                close(true);
                return false;
            }
            else if(s < 0)
            {
                auto e = errno;

                switch(e)
                {
                case EAGAIN:
                case ECONNRESET:
                    // we did not finish sending
                    return true;
                    break;
                case EPIPE:
                    lock.unlock();
                    close(true);
                    return false;
                default:
                    lock.unlock();
                    close(true);
                    throw socket_error(strerror(errno));
                }
            }
        }

        m_send_queue_size -= message.length;
        m_send_queue.erase(it);
        m_send_queue_cond.notify_all();
    }
}

} //namespace yael::network
<commit_msg>Check socket validity before queueing up more message<commit_after>#include "yael/network/TcpSocket.h"

#include <sstream>
#include <fcntl.h>
#include <sys/socket.h>
#include <sys/poll.h>
#include <sys/types.h>
#include <csignal>
#include <arpa/inet.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <iostream>
#include <cerrno>
#include <stdexcept>
#include <cstring>
#include <unistd.h>
#include <cassert>
#include <glog/logging.h>

#include "DatagramMessageSlicer.h"
#include "StreamMessageSlicer.h"

using namespace std;

namespace yael::network
{

constexpr int TRUE_FLAG = 1;

TcpSocket::TcpSocket(MessageMode mode, size_t max_send_queue_size)
    : m_port(0), m_is_ipv6(false), m_fd(-1), m_max_send_queue_size(max_send_queue_size)
{
    if(mode == MessageMode::Datagram)
    {
        m_slicer = std::make_unique<DatagramMessageSlicer>();
    }
    else if(mode == MessageMode::Stream)
    {
        m_slicer = std::make_unique<StreamMessageSlicer>();
    }
    else
    {
        throw std::runtime_error("Invalid message mode");
    }
}

TcpSocket::TcpSocket(MessageMode mode, int fd, size_t max_send_queue_size)
    : m_port(0), m_is_ipv6(false), m_fd(fd), m_max_send_queue_size(max_send_queue_size)
{
    if(mode == MessageMode::Datagram)
    {
        m_slicer = std::make_unique<DatagramMessageSlicer>();
    }
    else if(mode == MessageMode::Stream)
    {
        m_slicer = std::make_unique<StreamMessageSlicer>();
    }
    else
    {
        throw std::runtime_error("Invalid message mode");
    }

    uint32_t flags = fcntl(m_fd, F_GETFL, 0);
    flags = flags | O_NONBLOCK;
    fcntl(m_fd, F_SETFL, flags);

    update_port_number();
    calculate_remote_address();

    m_state = State::Connected;
}

TcpSocket::~TcpSocket()
{
    TcpSocket::close(true);
}

bool TcpSocket::has_messages() const
{
    return m_slicer->has_messages();
}

bool TcpSocket::create_fd()
{
    if(m_is_ipv6)
    {
        m_fd = socket(AF_INET6, SOCK_STREAM, IPPROTO_TCP);
    }
    else
    {
        m_fd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    }

    if(!is_valid())
    {
        return false;
    }

    ::setsockopt(m_fd, IPPROTO_TCP, TCP_NODELAY, reinterpret_cast<const char*>(&TRUE_FLAG), sizeof(TRUE_FLAG));

    return true;
}

void TcpSocket::update_port_number()
{
    if(m_is_ipv6)
    {
        struct sockaddr_in6 addr;
        socklen_t addrlen = sizeof(addr);
        getsockname(m_fd, reinterpret_cast<sockaddr*>(&addr), &addrlen);
        m_port = htons(addr.sin6_port);
    }
    else
    {
        struct sockaddr_in addr;
        socklen_t addrlen = sizeof(addr);
        getsockname(m_fd, reinterpret_cast<sockaddr*>(&addr), &addrlen);
        m_port = htons(addr.sin_port);
    }
}

bool TcpSocket::bind_socket(const Address& address)
{
    m_is_ipv6 = address.IPv6;

    if(!create_fd())
    {
        return false;
    }

    // Reuse address so we can quickly recover from crashes
    ::setsockopt(m_fd, SOL_SOCKET, SO_REUSEADDR, &TRUE_FLAG, sizeof(TRUE_FLAG));

    if(m_is_ipv6)
    {
        sockaddr_in6 sock_addr;
        address.get_sock_address6(sock_addr);

        if(::bind(m_fd, reinterpret_cast<sockaddr*>(&sock_addr), sizeof(sock_addr)) != 0)
        {
            return false;
        }
    }
    else
    {
        sockaddr_in sock_addr;
        address.get_sock_address(sock_addr);

        if(::bind(m_fd, reinterpret_cast<sockaddr*>(&sock_addr), sizeof(sock_addr)) != 0)
        {
            return false;
        }
   }

    update_port_number();
    return true;
}

bool TcpSocket::listen(const Address& address, uint32_t backlog)
{
    if(!bind_socket(address))
    {
        throw socket_error("Failed to bind socket!");
    }

    uint32_t flags = fcntl(m_fd, F_GETFL, 0);
    flags = flags | O_NONBLOCK;
    fcntl(m_fd, F_SETFL, flags);

    if(::listen(m_fd, backlog) == 0)
    {
        m_state = State::Listening;
        return true;
    }
    else
    {
        return false;
    }
}

uint16_t TcpSocket::port() const
{
    if(!is_valid())
    {
        throw socket_error("Cannot get port of non-existent socket");
    }

    return m_port;
}

bool TcpSocket::connect(const Address& address, const std::string& name)
{
    if(address.PortNumber == 0)
    {
        throw std::invalid_argument("Need to specify a port number");
    }

    if(name.empty())
    {
        m_is_ipv6 = address.IPv6;
        if(!create_fd())
        {
            throw socket_error(strerror(errno));
        }
    }
    else
    {
        Address my_addr = resolve_URL(name, ANY_PORT, address.IPv6);

        if(!bind_socket(my_addr))
        {
            throw socket_error(strerror(errno));
        }
    }

    // Set it blocking just for connect
    uint32_t flags = fcntl(m_fd, F_GETFL, 0);
    flags = flags & ~O_NONBLOCK;
    fcntl(m_fd, F_SETFL, flags);

    if(m_is_ipv6)
    {
        sockaddr_in6 sock_addr;
        address.get_sock_address6(sock_addr);

        if(::connect(m_fd, reinterpret_cast<const sockaddr*>(&sock_addr), sizeof(sock_addr)) != 0)
        {
            close();
            return false;
        }
    }
    else
    {
        sockaddr_in sock_addr;
        address.get_sock_address(sock_addr);

        if(::connect(m_fd, reinterpret_cast<const sockaddr*>(&sock_addr), sizeof(sock_addr)) != 0)
        {
            close();
            return false;
        }
    }

    flags = fcntl(m_fd, F_GETFL, 0);
    flags = flags | O_NONBLOCK;
    fcntl(m_fd, F_SETFL, flags);

    calculate_remote_address();
    update_port_number();

    m_state = State::Connected;
    return true;
}

int32_t TcpSocket::internal_accept()
{
    int32_t s = 0;

    if(m_is_ipv6)
    {
        sockaddr_in6 sock_addr;
        socklen_t len = sizeof(sock_addr);
        s = ::accept(m_fd, reinterpret_cast<sockaddr*>(&sock_addr), &len);
    }
    else
    {
        sockaddr_in sock_addr;
        socklen_t len = sizeof(sock_addr);
        s = ::accept(m_fd, reinterpret_cast<sockaddr*>(&sock_addr), &len);
    }

    if(s < 0 && errno != EWOULDBLOCK)
    {
        close();
        std::string str = "Failed to accept new connection; ";
        str += strerror(errno);
        throw socket_error(str);
    }

    return s;
}

std::vector<std::unique_ptr<Socket>> TcpSocket::accept()
{
    if(!is_listening())
    {
        throw socket_error("Cannot accept on connected TcpTcpSocket");
    }

    std::vector<std::unique_ptr<Socket>> res;

    while(true)
    {
        auto fd = internal_accept();

        if(fd >= 0)
        {
            auto ptr = new TcpSocket(m_slicer->type(), fd, m_max_send_queue_size);
            res.emplace_back(std::unique_ptr<Socket>(ptr));
        }
        else
        {
            return res;
        }
    }
}

const Address& TcpSocket::get_remote_address() const
{
    return m_remote_address;
}

void TcpSocket::calculate_remote_address()
{
    char ipstring[16];
    sockaddr_in sin;
    uint32_t len = sizeof(sin);

    if( getpeername(m_fd, reinterpret_cast<sockaddr*>(&sin), &len) == -1)
    {
        m_remote_address.reset();
    }

    inet_ntop( AF_INET, dynamic_cast<in_addr*>(&sin.sin_addr), &ipstring[0], 16);

    uint16_t port = htons(sin.sin_port);

    m_remote_address.IP = &ipstring[0];
    m_remote_address.PortNumber = port;
}

bool TcpSocket::close(bool fast)
{
    if(m_state == State::Connected && !fast)
    {
        m_state = State::Shutdown;
        int i = ::shutdown(m_fd, SHUT_RD | SHUT_WR);
        (void)i; //unused

        return false;
    }

    if(m_fd > 0)
    {
        m_state = State::Closed;
        int i = ::close(m_fd);
        (void)i; //unused
        m_fd = -1;

        m_slicer->buffer().reset();
    }
    else
    {
        if(!(m_state == State::Closed || m_state == State::Unknown))
        {
            LOG(FATAL) << "Invalid state";
        }

        //no-op
    }

    // threads might wait for send queue to be empty
    {
        std::unique_lock lock(m_send_mutex);
        m_send_queue_cond.notify_all();
    }

    return true;
}

void TcpSocket::pull_messages()
{
    auto &buffer = m_slicer->buffer();

    if(!buffer.is_valid())
    {
        bool res = receive_data(buffer);
        if(!res)
        {
            return;
        }
    }

    try
    {
        m_slicer->process_buffer();
    }
    catch(std::exception &e)
    {
        // ignore
    }

    // read rest of buffer
    // always pull more until we get EAGAIN
    pull_messages();
}

bool TcpSocket::receive_data(buffer_t &buffer)
{
    if(!is_valid())
    {
        return false;
    }

    if(buffer.is_valid())
    {
        throw std::runtime_error("TcpSocket::receive_data failed: Still have data queued up in buffer");
    }

    memset(&buffer.data[0], 0, yael::network::buffer_t::MAX_SIZE);
    auto x = ::recv(m_fd, buffer.data, yael::network::buffer_t::MAX_SIZE, 0);

    // Now act accordingly
    // > 0 -> data
    // = 0 -> disconnect
    // < 0 -> error/block
    if(x > 0)
    {
        buffer.size = x;
        buffer.position = 0;

        return true;
    }
    else if(x == 0)
    {
        close(true);
        return false;
    }
    else
    {
        const int e = errno;

        switch(e)
        {
        case EAGAIN:
            break;
        case ECONNRESET:
            close(true);
            break;
        default:
        {
            std::string str = "Failed to receive data; ";
            str += strerror(errno);

            // First close socket and then throw the error!
            close(true);
            throw socket_error(str);
        }
        }

        return false;
    }
}

std::optional<TcpSocket::message_in_t> TcpSocket::receive()
{
    pull_messages();

    if(m_slicer->has_messages())
    {
        message_in_t msg;
        auto res = m_slicer->get_message(msg);
        if(!res)
        {
            throw socket_error("failed to get message");
        }

        return { msg };
    }
    else
    {
        return {};
    }
}

bool TcpSocket::send(std::unique_ptr<uint8_t[]> &&data, uint32_t len)
{
    if(len <= 0)
    {
        throw socket_error("Message size has to be > 0");
    }

    if(!is_valid())
    {
        throw socket_error("Socket is closed");
    }

    m_slicer->prepare_message(data, len);

    auto msg_out = message_out_internal_t(std::move(data), len);

    {
        std::unique_lock lock(m_send_mutex);

        if(m_send_queue_size >= m_max_send_queue_size)
        {
            throw send_queue_full();
        }

        m_send_queue_size += msg_out.length;
        m_send_queue.emplace_back(std::move(msg_out));
    }

    return do_send();
}

void TcpSocket::wait_send_queue_empty()
{
    std::unique_lock lock(m_send_mutex);

    while(m_send_queue_size > 0 && is_valid())
    {
        m_send_queue_cond.wait(lock);
    }
}

bool TcpSocket::do_send()
{
    std::unique_lock lock(m_send_mutex);

    while(true)
    {
        auto it = m_send_queue.begin();

        if(it == m_send_queue.end())
        {
            // we sent everything!
            return false;
        }

        if(!is_valid())
        {
            throw socket_error("Socket is closed");
        }

        auto &message = *it;

        const uint32_t length = message.length;
        const uint8_t *rdata = message.data.get();

        while(message.sent_pos < length)
        {
            auto s = ::write(m_fd, rdata + message.sent_pos, length - message.sent_pos);

            if(s > 0)
            {
                message.sent_pos += s;
            }
            else if(s == 0)
            {
                LOG(WARNING) << "Connection lost during send: Message may only be sent partially";
                close(true);
                return false;
            }
            else if(s < 0)
            {
                auto e = errno;

                switch(e)
                {
                case EAGAIN:
                case ECONNRESET:
                    // we did not finish sending
                    return true;
                    break;
                case EPIPE:
                    lock.unlock();
                    close(true);
                    return false;
                default:
                    lock.unlock();
                    close(true);
                    throw socket_error(strerror(errno));
                }
            }
        }

        m_send_queue_size -= message.length;
        m_send_queue.erase(it);
        m_send_queue_cond.notify_all();
    }
}

} //namespace yael::network
<|endoftext|>


<commit_before>03813487-ad5a-11e7-b635-ac87a332f658<commit_msg>We apologise for the previous fix. Those responsible have been sacked<commit_after>040ac959-ad5a-11e7-8572-ac87a332f658<|endoftext|>


<commit_before>/*
 * Copyright (c) 2011-2013 libbitcoin developers (see AUTHORS)
 *
 * This file is part of libbitcoin.
 *
 * libbitcoin is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License with
 * additional permissions to the one published by the Free Software
 * Foundation, either version 3 of the License, or (at your option)
 * any later version. For more information see LICENSE.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program. If not, see <http://www.gnu.org/licenses/>.
 */
#include <bitcoin/bitcoin/network/handshake.hpp>

#include <cstdint>
#include <functional>
#include <boost/regex.hpp>
#include <boost/lexical_cast.hpp>
#include <bitcoin/bitcoin/constants.hpp>
#include <bitcoin/bitcoin/define.hpp>
#include <bitcoin/bitcoin/network/channel.hpp>
#include <bitcoin/bitcoin/network/network.hpp>
#include <bitcoin/bitcoin/primitives.hpp>
#include <bitcoin/bitcoin/utility/async_parallel.hpp>
#include <bitcoin/bitcoin/version.hpp>

namespace libbitcoin {
namespace network {

using std::placeholders::_1;
using std::placeholders::_2;

#define BC_USER_AGENT "/libbitcoin:" LIBBITCOIN_VERSION "/"

handshake::handshake(threadpool& pool, uint16_t port, uint32_t start_height)
  : strand_(pool.service())
{
    // TODO: shouldn't the nonce change on every handshake (like timestamp)?
    template_version_.nonce = rand();
    // template_version_.timestamp = time(nullptr);

    // Set fixed values inversion template.
    template_version_.version = bc::protocol_version;
    template_version_.user_agent = BC_USER_AGENT;
    template_version_.services = 1;

    // Set default values inversion template.
    template_version_.start_height = start_height;
    template_version_.address_me.services = template_version_.services;
    template_version_.address_me.ip = bc::localhost_ip_address;
    template_version_.address_me.port = port;
    template_version_.address_you.services = template_version_.services;
    template_version_.address_you.ip = bc::localhost_ip_address;
    template_version_.address_you.port = port;
}

void handshake::start(start_handler handle_start)
{
    discover_external_ip(std::bind(handle_start, _1));
}

void handshake::ready(channel_ptr node,
    handshake::handshake_handler handle_handshake)
{
    constexpr size_t sync = 3;

    // synchrnize three code paths (or error) before calling handle_handshake.
    const auto completion_callback = async_parallel(handle_handshake, sync);

    // Copy the version template and set its timestamp.
    auto session_version = template_version_;
    session_version.timestamp = time(nullptr);

    // TODO: where does session_version get customized?
    // Since we removed cURL discover_external_ip always returns localhost.
    // The port value was formerly hardwired to bc::protocol_port.

    node->send(session_version,
        strand_.wrap(std::bind(&handshake::handle_message_sent,
            this, _1, completion_callback)));

    node->subscribe_version(
        strand_.wrap(std::bind(&handshake::receive_version,
            this, _1, _2, node, completion_callback)));

    node->subscribe_verack(
        strand_.wrap(std::bind(&handshake::receive_verack,
            this, _1, _2, completion_callback)));
}

void handshake::handle_message_sent(const std::error_code& ec,
    handshake::handshake_handler completion_callback)
{
    completion_callback(ec);
}

void handshake::receive_version(const std::error_code& ec, const version_type&,
    channel_ptr node, handshake::handshake_handler completion_callback)
{
    if (ec)
        completion_callback(ec);
    else
        node->send(verack_type(),
            strand_.wrap(std::bind(&handshake::handle_message_sent,
                this, _1, completion_callback)));
}

void handshake::receive_verack(const std::error_code& ec, const verack_type&,
    handshake::handshake_handler completion_callback)
{
    completion_callback(ec);
}

void handshake::discover_external_ip(discover_ip_handler handle_discover)
{
    strand_.post(
        std::bind(&handshake::do_discover_external_ip,
            this, handle_discover));
}

void handshake::do_discover_external_ip(discover_ip_handler handle_discover)
{
    template_version_.address_me.ip = localhost_ip_address;
    handle_discover(std::error_code(), template_version_.address_me.ip);
}

void handshake::fetch_network_address(
    fetch_network_address_handler handle_fetch)
{
    strand_.post(
        std::bind(&handshake::do_fetch_network_address,
            this, handle_fetch));
}

void handshake::do_fetch_network_address(
    fetch_network_address_handler handle_fetch)
{
    handle_fetch(std::error_code(), template_version_.address_me);
}

void handshake::set_port(uint16_t port, setter_handler handle_set)
{
    strand_.post(
        std::bind(&handshake::do_set_port,
            this, port, handle_set));
}

void handshake::do_set_port(uint16_t port, setter_handler handle_set)
{
    template_version_.address_me.port = port;
    handle_set(std::error_code());
}

// TODO: deprecate (any reason to set this dynamically)?
void handshake::set_user_agent(const std::string& user_agent,
    setter_handler handle_set)
{
    strand_.post(
        std::bind(&handshake::do_set_user_agent,
            this, user_agent, handle_set));
}

// TODO: deprecate (any reason to set this dynamically)?
void handshake::do_set_user_agent(const std::string& user_agent,
    setter_handler handle_set)
{
    template_version_.user_agent = user_agent;
    handle_set(std::error_code());
}

void handshake::set_start_height(uint64_t height, setter_handler handle_set)
{
    strand_.post(
        std::bind(&handshake::do_set_start_height,
            this, height, handle_set));
}

void handshake::do_set_start_height(uint64_t height, setter_handler handle_set)
{
    // We type this method as uint64_t because that is what is returned by
    // fetch_last_height, whcih feeds directly into this method. But start_height
    // is uint32_t in the satoshi network protocol.
    BITCOIN_ASSERT(height <= bc::max_uint32);
    template_version_.start_height = static_cast<uint32_t>(height);
    handle_set(std::error_code());
}

void finish_connect(const std::error_code& ec, channel_ptr node,
    handshake& shake, network::connect_handler handle_connect)
{
    if (ec)
        handle_connect(ec, node);
    else
        shake.ready(node, std::bind(handle_connect, _1, node));
}

void connect(handshake& shake, network& net, const std::string& hostname,
    uint16_t port, network::connect_handler handle_connect)
{
    net.connect(hostname, port,
        std::bind(finish_connect,
            _1, _2, std::ref(shake), handle_connect));
}

} // namespace network
} // namespace libbitcoin

<commit_msg>Remove dead code.<commit_after>/*
 * Copyright (c) 2011-2013 libbitcoin developers (see AUTHORS)
 *
 * This file is part of libbitcoin.
 *
 * libbitcoin is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License with
 * additional permissions to the one published by the Free Software
 * Foundation, either version 3 of the License, or (at your option)
 * any later version. For more information see LICENSE.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program. If not, see <http://www.gnu.org/licenses/>.
 */
#include <bitcoin/bitcoin/network/handshake.hpp>

#include <cstdint>
#include <functional>
#include <boost/lexical_cast.hpp>
#include <bitcoin/bitcoin/constants.hpp>
#include <bitcoin/bitcoin/define.hpp>
#include <bitcoin/bitcoin/network/channel.hpp>
#include <bitcoin/bitcoin/network/network.hpp>
#include <bitcoin/bitcoin/primitives.hpp>
#include <bitcoin/bitcoin/utility/async_parallel.hpp>
#include <bitcoin/bitcoin/version.hpp>

namespace libbitcoin {
namespace network {

using std::placeholders::_1;
using std::placeholders::_2;

#define BC_USER_AGENT "/libbitcoin:" LIBBITCOIN_VERSION "/"

handshake::handshake(threadpool& pool, uint16_t port, uint32_t start_height)
  : strand_(pool.service())
{
    // template_version_.nonce = rand();
    // template_version_.timestamp = time(nullptr);

    // Set fixed values inversion template.
    template_version_.version = bc::protocol_version;
    template_version_.user_agent = BC_USER_AGENT;
    template_version_.services = 1;

    // Set default values inversion template.
    template_version_.start_height = start_height;
    template_version_.address_me.services = template_version_.services;
    template_version_.address_me.ip = bc::localhost_ip_address;
    template_version_.address_me.port = port;
    template_version_.address_you.services = template_version_.services;
    template_version_.address_you.ip = bc::localhost_ip_address;
    template_version_.address_you.port = port;
}

void handshake::start(start_handler handle_start)
{
    discover_external_ip(std::bind(handle_start, _1));
}

void handshake::ready(channel_ptr node,
    handshake::handshake_handler handle_handshake)
{
    constexpr size_t sync = 3;

    // synchrnize three code paths (or error) before calling handle_handshake.
    const auto completion_callback = async_parallel(handle_handshake, sync);

    // Copy the version template and set its timestamp.
    auto session_version = template_version_;
    template_version_.nonce = rand();
    session_version.timestamp = time(nullptr);

    // Since we removed cURL discover_external_ip always returns localhost.
    // The port value was formerly hardwired to bc::protocol_port.

    node->send(session_version,
        strand_.wrap(std::bind(&handshake::handle_message_sent,
            this, _1, completion_callback)));

    node->subscribe_version(
        strand_.wrap(std::bind(&handshake::receive_version,
            this, _1, _2, node, completion_callback)));

    node->subscribe_verack(
        strand_.wrap(std::bind(&handshake::receive_verack,
            this, _1, _2, completion_callback)));
}

void handshake::handle_message_sent(const std::error_code& ec,
    handshake::handshake_handler completion_callback)
{
    completion_callback(ec);
}

void handshake::receive_version(const std::error_code& ec, const version_type&,
    channel_ptr node, handshake::handshake_handler completion_callback)
{
    if (ec)
        completion_callback(ec);
    else
        node->send(verack_type(),
            strand_.wrap(std::bind(&handshake::handle_message_sent,
                this, _1, completion_callback)));
}

void handshake::receive_verack(const std::error_code& ec, const verack_type&,
    handshake::handshake_handler completion_callback)
{
    completion_callback(ec);
}

void handshake::discover_external_ip(discover_ip_handler handle_discover)
{
    strand_.post(
        std::bind(&handshake::do_discover_external_ip,
            this, handle_discover));
}

void handshake::do_discover_external_ip(discover_ip_handler handle_discover)
{
    template_version_.address_me.ip = localhost_ip_address;
    handle_discover(std::error_code(), template_version_.address_me.ip);
}

void handshake::fetch_network_address(
    fetch_network_address_handler handle_fetch)
{
    strand_.post(
        std::bind(&handshake::do_fetch_network_address,
            this, handle_fetch));
}

void handshake::do_fetch_network_address(
    fetch_network_address_handler handle_fetch)
{
    handle_fetch(std::error_code(), template_version_.address_me);
}

void handshake::set_port(uint16_t port, setter_handler handle_set)
{
    strand_.post(
        std::bind(&handshake::do_set_port,
            this, port, handle_set));
}

void handshake::do_set_port(uint16_t port, setter_handler handle_set)
{
    template_version_.address_me.port = port;
    handle_set(std::error_code());
}

// TODO: deprecate (any reason to set this dynamically)?
void handshake::set_user_agent(const std::string& user_agent,
    setter_handler handle_set)
{
    strand_.post(
        std::bind(&handshake::do_set_user_agent,
            this, user_agent, handle_set));
}

// TODO: deprecate (any reason to set this dynamically)?
void handshake::do_set_user_agent(const std::string& user_agent,
    setter_handler handle_set)
{
    template_version_.user_agent = user_agent;
    handle_set(std::error_code());
}

void handshake::set_start_height(uint64_t height, setter_handler handle_set)
{
    strand_.post(
        std::bind(&handshake::do_set_start_height,
            this, height, handle_set));
}

void handshake::do_set_start_height(uint64_t height, setter_handler handle_set)
{
    // We type this method as uint64_t because that is what is returned by
    // fetch_last_height, whcih feeds directly into this method. But start_height
    // is uint32_t in the satoshi network protocol.
    BITCOIN_ASSERT(height <= bc::max_uint32);
    template_version_.start_height = static_cast<uint32_t>(height);
    handle_set(std::error_code());
}

void finish_connect(const std::error_code& ec, channel_ptr node,
    handshake& shake, network::connect_handler handle_connect)
{
    if (ec)
        handle_connect(ec, node);
    else
        shake.ready(node, std::bind(handle_connect, _1, node));
}

void connect(handshake& shake, network& net, const std::string& hostname,
    uint16_t port, network::connect_handler handle_connect)
{
    net.connect(hostname, port,
        std::bind(finish_connect,
            _1, _2, std::ref(shake), handle_connect));
}

} // namespace network
} // namespace libbitcoin

<|endoftext|>